Awesome

🌄 Efficient Test-Time Model Adaptation without Forgetting

This is the official project repository for Efficient Test-Time Model Adaptation without Forgetting 🔗 by Shuaicheng Niu, Jiaxiang Wu, Yifan Zhang, Yaofo Chen, Shijian Zheng, Peilin Zhao and Mingkui Tan (ICML 2022).

🌄 EATA 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.

• 1️⃣: EATA conducts selective/sample-adaptative optimization, i.e., only perform backward propagation for reliable and non-redundant test samples, which improves adaptaion efficiency and performance simtaneously.

• 2️⃣: EATA regularizes the model parameters during testing to prevent the forgetting on in-distribution/clean test samples 😋.

Installation:

EATA depends on

• Python 3
• PyTorch >= 1.0

Data preparation:

This repository contains code for evaluation on ImageNet and ImageNet-C 🔗 with ResNet models. But feel free to use your own data and models!

• Step 1: Download ImageNet-C 🔗 dataset from here 🔗.

• Step 2: Put IamgeNet-C at "--data_corruption" and put ImageNet test/val set at "--data".

Usage:

import eata

model = TODO_model()

model = eata.configure_model(model)
params, param_names = eata.collect_params(model)
optimizer = TODO_optimizer(params, lr=2.5e-4)
adapt_model = eata.EATA(model, optimizer, fishers, e_margin, d_margin) 

outputs = adapt_model(inputs)  # now it infers and adapts!

Notes:

• fishers are pre-calculated (see main.py) for preventing forgetting,
• e_margin and d_margin are two parameters for selective (sample-adaptive) optimization.

Example: Adapting a pre-trained ResNet-50 model on ImageNet-C (Corruption).

Usage:

python3 main.py --data /path/to/imagenet --data_corruption /path/to/imagenet-c --exp_type 'continual' or 'each_shift_reset' --algorithm 'eata' or 'eta' or 'tent' --output /output/dir

'--exp_type' is choosen from:

• 'continual' means the model parameters will never be reset, also called online adaptation;

• 'each_shift_reset' means after each type of distribution shift, e.g., ImageNet-C Gaussian Noise Level 5, the model parameters will be reset.

'--algorithm' is chosen from:

• 'tent' (baseline);
• 'eta' (ours eata w/o regularization);
• 'eata' (ours)

Results:

Here, we report the results on ImageNet-C, severity level = 5, with ResNet-50.

• Corruption accuracy (ETA/EATA achieves higher corruption accuracy but use fewer backward passes (compared to TTT, Tent), thereby more efficient):

• Clean accuracy (testing the model's source accuracy on clean/original imagenet test set). EATA improves model's corruption acc. and maintains the source acc., while Tent can not.

Please see our PAPER 🔗 for detailed results.

Correspondence

Please contact Shuaicheng Niu by niushuaicheng [at] gmail.com 📬.

Citation

If the EATA method or fully test-time adaptation without forgetting are helpful in your research, please consider citing our paper:

@InProceedings{niu2022efficient,
  title={Efficient Test-Time Model Adaptation without Forgetting},
  author={Niu, Shuaicheng and Wu, Jiaxiang and Zhang, Yifan and Chen, Yaofo and Zheng, Shijian and Zhao, Peilin and Tan, Mingkui},
  booktitle = {The Internetional Conference on Machine Learning},
  year = {2022}
}

Acknowledgment

The code is greatly inspired by (heavily from) the Tent 🔗 and TTT 🔗.