Test-Time Classifier Adjustment Module for Model-Agnostic Domain Generalization
This codebase is the official implementation of Test-Time Classifier Adjustment Module for Model-Agnostic Domain Generalization (NeurIPS2021, Spotlight).
This codebase is mainly based on DomainBed, with following modifications:
- enable to use various backbone networks including Big Transfer (BiT), Vision Transformers (ViT, DeiT, HViT), and MLP-Mixer.
- enable to test test-time adaptation method (T3A and Tent).
Installation
CUDA/Python
git clone [email protected]:matsuolab/Domainbed_contrib.git
cd Domainbed_contrib/docker
docker build -t {image_name} .
docker run -it -h `hostname` --runtime=nvidia -v /path/to/Domainbed_contrib /path/to/anyware --shm-size=40gb --name {container_name} {image_name}
Python libralies
We use pipenv for package management.
cd /path/to/Domainbed_contrib
pip install pipenv
pipenv install
pipenv shell
pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.8.0+cu102.html
Quick start
(1) Downlload the datasets
python -m domainbed.scripts.download --data_dir=/my/datasets/path --dataset pacs
Note: change --dataset pacs for downloading other datasets (e.g., vlcs, office_home, terra_incognita).
(2) Train a model on source domains
python -m domainbed.scripts.train\
--data_dir /my/datasets/path\
--output_dir /my/pretrain/path\
--algorithm ERM\
--dataset PACS\
--hparams "{\"backbone\": \"resnet50\"}"
This scripts will produce new directory /my/pretrain/path, which include the full training log.
Note: change --dataset PACS for training on other datasets (e.g., VLCS, OfficeHome, TerraIncognita).
Note: change --hparams "{\"backbone\": \"resnet50\"}" for using other backbones (e.g., resnet18, ViT-B16, HViT).
(3) Evaluate model with test time adaptation (Table 1, Table 2, Figure 2)
python -m domainbed.scripts.unsupervised_adaptation\
--input_dir=/my/pretrain/path\
--adapt_algorithm=T3A
This scripts will produce a new file in /my/pretrain/path, whose name is results_{adapt_algorithm}.jsonl.
Note: change --adapt_algorithm=T3A for using other test time adaptation methods (T3A, Tent, or TentClf).
(4) Evaluate model with fine-tuning classifier(Figure 1)
python -m domainbed.scripts.supervised_adaptation\
--input_dir=/my/pretrain/path\
--ft_mode=clf
This scripts will produce a new file in /my/pretrain/path, whose name is results_{ft_mode}.jsonl.
Available backbones
- resnet18
- resnet50
- BiT-M-R50x3
- BiT-M-R101x3
- BiT-M-R152x2
- ViT-B16
- ViT-L16
- DeiT
- Hybrid ViT (HViT)
- MLP-Mixer (Mixer-L16)
Reproducing results
Table 1 and Figure 2 (Tuned ERM and CORAL)
You can use scripts/hparam_search.sh. Specifically, for each dataset and base algorithm, you can just type a following command.
sh scripts/hparam_search.sh resnet50 PACS ERM
Note that, it automatically starts 240 jobs, and take many times to finish.
Table 2 and Figure 1 (ERM with various backbone)
You can use scripts/launch.sh. Specifically, for each backbone, you can just type following three commands.
sh scripts/launch.sh pretrain resnet50 10 3 local
sh scripts/launch.sh sup resnet50 10 3 local
sh scripts/launch.sh unsup resnet50 10 3 local
Other results
For table 1, we used scores reported by In Search of Lost Domain Generalization.
Full results for the reported scores in LaTeX format available here.
Note: We only used scores for VLCS, PACS, OfficeHome, and TerraIncognita. We used the resutls with IIDAccuracySelectionMethod.
License
This source code is released under the MIT license, included here.
