Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNet, ICCV 2021
Update:
2021/03/11: update our new results. Now our T2T-ViT-14 with 21.5M parameters can reach 81.5% top1-acc with 224x224 image resolution, and 83.3% top1-acc with 384x384 resolution.
2021/02/21: T2T-ViT can be trained on most of common GPUs: 1080Ti, 2080Ti, TiTAN V, V100 stably with '--amp' (Automatic Mixed Precision). In some specifical GPU like Tesla T4, 'amp' would cause NAN loss when training T2T-ViT. If you get NAN loss in training, you can disable amp by removing '--amp' in the training scripts.
2021/01/28: release codes and upload most of the pretrained models of T2T-ViT.
Reference
If you find this repo useful, please consider citing:
@InProceedings{Yuan_2021_ICCV,
author = {Yuan, Li and Chen, Yunpeng and Wang, Tao and Yu, Weihao and Shi, Yujun and Jiang, Zi-Hang and Tay, Francis E.H. and Feng, Jiashi and Yan, Shuicheng},
title = {Tokens-to-Token ViT: Training Vision Transformers From Scratch on ImageNet},
booktitle = {Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV)},
month = {October},
year = {2021},
pages = {558-567}
}
Our codes are based on the official imagenet example by PyTorch and pytorch-image-models by Ross Wightman
1. Requirements
timm, pip install timm==0.3.4
torch>=1.4.0
torchvision>=0.5.0
pyyaml
data prepare: ImageNet with the following folder structure, you can extract imagenet by this script.
│imagenet/
├──train/
│ ├── n01440764
│ │ ├── n01440764_10026.JPEG
│ │ ├── n01440764_10027.JPEG
│ │ ├── ......
│ ├── ......
├──val/
│ ├── n01440764
│ │ ├── ILSVRC2012_val_00000293.JPEG
│ │ ├── ILSVRC2012_val_00002138.JPEG
│ │ ├── ......
│ ├── ......
2. T2T-ViT Models
| Model | T2T Transformer | Top1 Acc | #params | MACs | Download |
|---|---|---|---|---|---|
| T2T-ViT-14 | Performer | 81.5 | 21.5M | 4.8G | here |
| T2T-ViT-19 | Performer | 81.9 | 39.2M | 8.5G | here |
| T2T-ViT-24 | Performer | 82.3 | 64.1M | 13.8G | here |
| T2T-ViT-14, 384 | Performer | 83.3 | 21.7M | here | |
| T2T-ViT-24, Token Labeling | Performer | 84.2 | 65M | here | |
| T2T-ViT_t-14 | Transformer | 81.7 | 21.5M | 6.1G | here |
| T2T-ViT_t-19 | Transformer | 82.4 | 39.2M | 9.8G | here |
| T2T-ViT_t-24 | Transformer | 82.6 | 64.1M | 15.0G | here |
The 'T2T-ViT-14, 384' means we train T2T-ViT-14 with image size of 384 x 384.
The 'T2T-ViT-24, Token Labeling' means we train T2T-ViT-24 with Token Labeling.
The three lite variants of T2T-ViT (Comparing with MobileNets):
| Model | T2T Transformer | Top1 Acc | #params | MACs | Download |
|---|---|---|---|---|---|
| T2T-ViT-7 | Performer | 71.7 | 4.3M | 1.1G | here |
| T2T-ViT-10 | Performer | 75.2 | 5.9M | 1.5G | here |
| T2T-ViT-12 | Performer | 76.5 | 6.9M | 1.8G | here |
Usage
The way to use our pretrained T2T-ViT:
from models.t2t_vit import *
from utils import load_for_transfer_learning
# create model
model = t2t_vit_14()
# load the pretrained weights
load_for_transfer_learning(model, /path/to/pretrained/weights, use_ema=True, strict=False, num_classes=1000) # change num_classes based on dataset, can work for different image size as we interpolate the position embeding for different image size.
3. Validation
Test the T2T-ViT-14 (take Performer in T2T module),
Download the T2T-ViT-14, then test it by running:
CUDA_VISIBLE_DEVICES=0 python main.py path/to/data --model t2t_vit_14 -b 100 --eval_checkpoint path/to/checkpoint
The results look like:
Test: [ 0/499] Time: 2.083 (2.083) Loss: 0.3578 (0.3578) [email protected]: 96.0000 (96.0000) [email protected]: 99.0000 (99.0000)
Test: [ 50/499] Time: 0.166 (0.202) Loss: 0.5823 (0.6404) [email protected]: 85.0000 (86.1569) [email protected]: 99.0000 (97.5098)
...
Test: [ 499/499] Time: 0.272 (0.172) Loss: 1.3983 (0.8261) [email protected]: 62.0000 (81.5000) [email protected]: 93.0000 (95.6660)
Top-1 accuracy of the model is: 81.5%
Test the three lite variants: T2T-ViT-7, T2T-ViT-10, T2T-ViT-12 (take Performer in T2T module),
Download the T2T-ViT-7, T2T-ViT-10 or T2T-ViT-12, then test it by running:
CUDA_VISIBLE_DEVICES=0 python main.py path/to/data --model t2t_vit_7 -b 100 --eval_checkpoint path/to/checkpoint
Test the model T2T-ViT-14, 384 with 83.3% top-1 accuracy:
CUDA_VISIBLE_DEVICES=0 python main.py path/to/data --model t2t_vit_14 --img-size 384 -b 100 --eval_checkpoint path/to/T2T-ViT-14-384
4. Train
Train the three lite variants: T2T-ViT-7, T2T-ViT-10 and T2T-ViT-12 (take Performer in T2T module):
If only 4 GPUs are available,
CUDA_VISIBLE_DEVICES=0,1,2,3 ./distributed_train.sh 4 path/to/data --model t2t_vit_7 -b 128 --lr 1e-3 --weight-decay .03 --amp --img-size 224
The top1-acc in 4 GPUs would be slightly lower than 8 GPUs (around 0.1%-0.3% lower).
If 8 GPUs are available:
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 path/to/data --model t2t_vit_7 -b 64 --lr 1e-3 --weight-decay .03 --amp --img-size 224
Train the T2T-ViT-14 and T2T-ViT_t-14 (run on 4 or 8 GPUs):
CUDA_VISIBLE_DEVICES=0,1,2,3 ./distributed_train.sh 4 path/to/data --model t2t_vit_14 -b 128 --lr 1e-3 --weight-decay .05 --amp --img-size 224
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 path/to/data --model t2t_vit_14 -b 64 --lr 5e-4 --weight-decay .05 --amp --img-size 224
If you want to train our T2T-ViT on images with 384x384 resolution, please use '--img-size 384'.
Train the T2T-ViT-19, T2T-ViT-24 or T2T-ViT_t-19, T2T-ViT_t-24:
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 path/to/data --model t2t_vit_19 -b 64 --lr 5e-4 --weight-decay .065 --amp --img-size 224
5. Transfer T2T-ViT to CIFAR10/CIFAR100
| Model | ImageNet | CIFAR10 | CIFAR100 | #params |
|---|---|---|---|---|
| T2T-ViT-14 | 81.5 | 98.3 | 88.4 | 21.5M |
| T2T-ViT-19 | 81.9 | 98.4 | 89.0 | 39.2M |
We resize CIFAR10/100 to 224x224 and finetune our pretrained T2T-ViT-14/19 to CIFAR10/100 by running:
CUDA_VISIBLE_DEVICES=0,1 transfer_learning.py --lr 0.05 --b 64 --num-classes 10 --img-size 224 --transfer-learning True --transfer-model /path/to/pretrained/T2T-ViT-19
6. Visualization
Visualize the image features of ResNet50, you can open and run the visualization_resnet.ipynb file in jupyter notebook or jupyter lab; some results are given as following:
Visualize the image features of ViT, you can open and run the visualization_vit.ipynb file in jupyter notebook or jupyter lab; some results are given as following:
Visualize attention map, you can refer to this file. A simple example by visualizing the attention map in attention block 4 and 5 is:
The 2nd train with Adam optimizer and cosine-warmup.（not use timm.create_optimizer to set adamw sice i need to set different lr for different blocks.） The set of lr is similar to your setting.After 40 epochs, the loss curves also as follow.
It's look like that the 2nd training much better and the loss is still in decrease.But I'm not sure is it on the right path.(according to my calculation, it will take 6 days to train 300 epochs with a single 3090 GPU, so I don't have time to trial & error:sob::sob::sob:)
Could you show me your training log as a reference or give me some advice? Thank you very much.
![[maybe a bug] loss nan](https://avatars.githubusercontent.com/u/43675899?v=4)
2 Feb 09, 2022
145 Jan 04, 2023
15 Nov 18, 2022
51 Dec 01, 2022
39 Dec 25, 2022
135 Oct 27, 2022
252 Sep 16, 2022
604 Dec 14, 2022
47 Nov 06, 2022
85 Jan 03, 2023
108 Dec 31, 2022
3k Dec 31, 2022
967 Jan 04, 2023
16.4k Jan 09, 2023
304 Dec 25, 2022
6 Dec 14, 2022
22 Oct 28, 2022
20 Dec 06, 2022
177 Jan 04, 2023
208 Dec 20, 2022