Improving Model Generalization by Agreement of Learned Representations from Data Augmentation (WACV 2022)
Paper
Why it matters?
When data augmentation is applied on an input image, a model is forced to learn invariant features to improve model generalization (Figure 1).
Since data augmentation incurs little overhead, why not generate 2 data augmented images (also known as 2 positive samples) from a given input. Then, force the model to agree on the common invariant features to support the correct label (Figure 2). It turns out that maximizing this agreement further improves model model generalization. We call our method AgMax.
Unlike label smoothing, AgMax consistently improves model accuracy. For example on ImageNet1k for 90 epochs, the ResNet50 performance is as follows:
| Data Augmentation | Baseline | Label Smoothing | AgMax (Ours) |
|---|---|---|---|
| Standard | 76.4 | 76.8 | 76.9 |
| CutOut | 76.2 | 76.5 | 77.1 |
| MixUp | 76.5 | 76.7 | 77.6 |
| CutMix | 76.3 | 76.4 | 77.4 |
| AutoAugment (AA) | 76.2 | 76.2 | 77.1 |
| CutOut+AA | 75.7 | 75.7 | 76.6 |
| MixUp+AA | 75.9 | 76.5 | 77.1 |
| CutMix+AA | 75.5 | 75.5 | 77.0 |
The figure below demonstrates consistent improvement across different data augmnentation methods:
Install requirements
pip3 install -r requirements.txt
Train
For example, train ResNet50 with AgMax on 2 GPUs for 90 epochs, SGD with lr=0.1 and multistep learning rate scheduler:
CUDA_VISIBLE_DEVICES=0,1 python3 main.py --config=ResNet50-standard-agmax --train \
--multisteplr --dataset=imagenet --epochs=90 --save
Compare the results without AgMax:
CUDA_VISIBLE_DEVICES=0,1 python3 main.py --config=ResNet50-standard --train \
--multisteplr --dataset=imagenet --epochs=90 --save
Test
Using a pre-trained model:
ResNet101 trained with CutMix, AutoAugment and AgMax:
mkdir checkpoints
cd checkpoints
wget https://github.com/roatienza/agmax/releases/download/agmax-0.1.0/imagenet-agmax-mi-ResNet101-cutmix-auto_augment-81.19-mlp-4096.pth
cd ..
python3 main.py --config=ResNet101-auto_augment-cutmix-agmax --eval \
--dataset=imagenet \
--resume imagenet-agmax-mi-ResNet101-cutmix-auto_augment-81.19-mlp-4096.pth
ResNet50 trained with CutMix, AutoAugment and AgMax:
python3 main.py --config=ResNet50-auto_augment-cutmix-agmax --eval --n-units=2048 \
--dataset=imagenet --resume imagenet-agmax-ResNet50-cutmix-auto_augment-79.12-mlp-2048.pth
Other pre-trained models (Baselines):
Citation
If you find this work useful, please cite:
@inproceedings{atienza2022agmax,
title={Improving Model Generalization by Agreement of Learned Representations from Data Augmentation},
author={Atienza, Rowel},
booktitle = {IEEE/CVF Winter Conference on Applications of Computer Vision},
year={2022},
pubstate={published},
tppubtype={inproceedings}
}
16 Dec 14, 2022
85 Jan 2, 2023
54 Dec 8, 2022
23 Oct 9, 2022
25 Sep 26, 2022
97 Dec 13, 2022
82 Jan 4, 2023
3 Aug 20, 2022
38 Nov 1, 2022
83 Dec 31, 2022
536 Jan 05, 2023
1.2k Jan 03, 2023
200 Dec 20, 2022
387 Jan 08, 2023
185 Jan 09, 2023
10 Oct 06, 2022
72 Dec 13, 2022
50 Sep 26, 2022
3 Sep 20, 2022
76 Sep 05, 2022
12 Nov 05, 2022
38 Dec 13, 2022
501 Jan 03, 2023
9 Dec 15, 2022
177 Dec 11, 2022
589 Dec 30, 2022
21 Nov 29, 2022
185 Jan 07, 2023