Fast-AutoAug-Torch
Search for Fast AutoAugment using AutoTorch. AutoAugment and RandAugment are also implemented for comparison.
This example and Search for RegNet will be used in the tutorial on From HPO to NAS: Automated Deep Learning at CVPR 2020.
| model | augment | epoch | Acc | weights |
|---|---|---|---|---|
| ResNet-50 | baseline | 120 | 76.48 | |
| ResNet-50 | AA | 120 | 76.66 | |
| ResNet-50 | Fast AA | 120 | 76.88 | |
| ResNet-50 | Rand AA | 120 | 76.79 | |
| ResNet-50 | baseline | 270 | 77.17 | |
| ResNet-50 | AA | 270 | 77.78 | link |
| ResNet-50 | Fast AA | 270 | 77.73 | link |
| ResNet-50 | Rand AA | 270 | 77.97 | link |
Approaches:
AA: AutoAugment, Fast AA: Fast AutoAugment, Rand AA: RandAugment,
Training HP setting:
learning rate: 0.2, batch size: 512, weight decay: 1e-4,
Fast AA setting (using random search):
K=8, T=1, N=5,
resulting 40 subpolicies in total.
RandAug setting (after grid search):
n=2, m=12,
Setup
Install dependencies
-
PyTorch, please follow the instructions.
-
Install AutoTorch and PyTorch-Encoding toolkits:
pip install autotorch --pre
pip install torch-encoding --pre- Install Apex (optional):
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
Prepare dataset
# assuming you have downloaded the ImageNet dataset in the current folder
python prepare_imagenet.py --download-dir ./Fast AutoAugment Training
Search policy
If you want to skip this step, you can download the searched policy here.
python search_policy.py --reduced-size 60000 --epochs 120 --nfolds 8 --num-trials 200 --save-policy imagenet_policy.atTrain with searched policy
python train.py --dataset imagenet --model resnet50 --lr-scheduler cos --epochs 270 --checkname resnet50_fast_aa --lr 0.025 --batch-size 64 --auto-policy imagenet_policy.atOther Methods
Training with AutoAugment policy
For AutoAugment we only enable training with searched policy from the original paper.
# download the policy
wget https://hangzh.s3-us-west-1.amazonaws.com/others/aa_policy.at
# start the training
python train.py --dataset imagenet --model resnet50 --lr-scheduler cos --epochs 270 --checkname resnet50_fast_aa --lr 0.025 --batch-size 64 --auto-policy aa_policy.atTraining with RandomAugment
python train.py --dataset imagenet --model resnet50 --lr-scheduler cos --epochs 270 --checkname resnet50_rand_aug --lr 0.025 --batch-size 64 --rand-aug