Test-Time Model Adaptation with Only Forward Passes

This is the official project repository for Test-Time Model Adaptation with Only Forward Passes (ICML 2024, Oral) by Shuaicheng Niu, Chunyan Miao, Guohao Chen, Pengcheng Wu, Peilin Zhao

1️⃣ FOA conducts model learning at test time to adapt a pre-trained model to test data that has distributional shifts ☀️ 🌧 ❄️, such as corruptions, simulation-to-real discrepancies, and other differences between training and testing data.
2️⃣ FOA performs adaptation on both input and output levels, which avoids modification to model parameters and adapts in a backpropagation-free manner. Consequently, FOA offers the following benefits:
1. reduces memory usage significantly, e.g., 5,165MB (Tent) $\rightarrow$ 832MB when using ViTBase and a BS of 64.
2. compatible with quantized models, which typically do not support backpropagation.
3. compatible with models on specialized chips, where parameters are hardcoded and non-modifiable.

Dependencies Installation:

pip install cma
pip install torch==2.1.0 torchvision==0.16.0 torchaudio==2.1.0 --index-url https://download.pytorch.org/whl/cu118
pip install timm==0.9.10

Data Preparation:

This repository contains code for evaluation on ImageNet-C/R/V2/Sketch with VitBase. But feel free to use your own data and models! Please check here 🔗 for a detailed guide on preparing these datasets.

Usage

from tta_library.foa import FOA
from models.vpt import PromptViT

model = TODO_model()
model = PromptViT(model, 3)

adapt_model = FOA(model, args.fitness_lambda)
train_loader = TODO_loader()
adapt_model.obtain_origin_stat(train_loader)

outputs = adapt_model(inputs)

Example: ImageNet-C Experiments

Usage (full precision experiments):

python3 main.py \
    --data path/to/imagenet \
    --data_v2 path/to/imagenet-v2 \
    --data_sketch path/to/imagenet-sketch \
    --data_corruption path/to/imagenet-c \
    --data_rendition path/to/imagenet-r \
    --algorithm [tent/foa/lame/t3a/sar/cotta] \

For experiments with quantized ViT, simply add

--quant

in the above command.

Experimental Results

The Table below demonstrates the result of both full precision and quantized ViTs. The reported average accuracy (%, $\uparrow$) / ECE (%, $\downarrow$) is averaged over 15 different corruption types in ImageNet-C (severity level 5).

	ViT (full precision, 32-bit)	ViT (8-bit)	ViT (6-bit)
NoAdapt	55.5 / 10.5	54.1 / 10.8	47.7 / 9.9
LAME	54.1 / 11.0	52.5 / 12.4	45.8 / 10.4
T3A	56.9 / 26.8	55.1 / 25.9	45.4 / 30.1
Tent	59.6 / 18.5	-	-
CoTTA	61.7 / 6.5	-	-
SAR	62.7 / 7.0	-	-
FOA	66.3 / 3.2	63.5 / 3.8	55.8 / 5.5

Please see our PAPER 🔗 for more detailed results.

Correspondence

Please contact Shuaicheng Niu by [shuaicheng.niu at ntu.edu.sg] and Guohao Chen by [chenguohao987 at gmail.com] if you have any questions. 📬

Citation

If our FOA method or the setting of test-time adaptation with only forward passes is helpful in your research, please consider citing our paper:

@inproceedings{niu2024test,
  title={Test-Time Model Adaptation with Only Forward Passes},
  author={Shuaicheng Niu, Chunyan Miao, Guohao Chen, Pengcheng Wu, Peilin Zhao},
  booktitle = {The International Conference on Machine Learning},
  year = {2024}
}

Acknowledgment

The code is inspired by EATA 🔗 and VPT 🔗

mr-eggplant/FOA

Test-Time Model Adaptation with Only Forward Passes

Example: ImageNet-C Experiments

Correspondence

Citation

Acknowledgment