Code release for Test-Time Training with Self-Supervision for Generalization under Distribution Shifts.
This code produces our results on ImageNet-C and ImageNet-Video-Robust.
The CIFAR-10 results are produced by this repository.
- Our code requires pytorch version 1.0 or higher, with at least one modern GPU of adequate memory.
- We ran our code with python 3.7. Compatibility with python 2 is possible maybe with some modifications.
- Most of the packages used should be included with anaconda, except maybe two small utilities:
- Download the two datasets into the same folder:
- ImageNet-C (Hendrycks and Dietterich) from this repository.
- ImageNet-Video-Robust (Shankar et al.) from this repository.
- Clone our repository with
git clone https://github.com/yueatsprograms/ttt_imagenet_release
. - Inside the repository, set the data folder to where the datasets are stored by editing:
--dataroot
argument inmain.py
.--dataroot
argument intest_video.py
.dataroot
variable inscript_test.py
.
- Run
script.sh
. - The results are stored in the respective folders in
results/
. - Once everything is finished, the results can be compiled and visualized with the following utilities:
show_table.py
parses the results into tables and prints them.show_plot.py
makes bar plots like those in our paper, and prints the tables in latex format; requires first runningshow_table.py
.