Package: tlt
All Tokens Matter: Token Labeling for Training Better Vision Transformers (arxiv)This is a Pytorch implementation of our paper.Comparison between the proposed LV-ViT and other recent works based on transformers. Note that we only show models whose model sizes are under 100M.Our codes are based on thepytorch-image-modelsbyRoss Wightman.Update2021.7: Add script to generate label data.2021.6: Supportpip install tltto use our Token Labeling Toolbox for image models.2021.6: Release training code and segmentation model.2021.4: Release LV-ViT models.LV-ViT ModelsModellayerdimImage resolutionParamTop 1DownloadLV-ViT-S1638422426.15M83.3linkLV-ViT-S1638438426.30M84.4linkLV-ViT-M2051222455.83M84.0linkLV-ViT-M2051238456.03M85.4linkLV-ViT-M2051244856.13M85.5linkLV-ViT-L24768448150.47M86.2linkLV-ViT-L24768512150.66M86.4linkRequirementstorch>=1.4.0
torchvision>=0.5.0
pyyaml
scipy
timm==0.4.5data prepare: ImageNet with the following folder structure, you can extract imagenet by thisscript.│imagenet/
├──train/
│  ├── n01440764
│  │   ├── n01440764_10026.JPEG
│  │   ├── n01440764_10027.JPEG
│  │   ├── ......
│  ├── ......
├──val/
│  ├── n01440764
│  │   ├── ILSVRC2012_val_00000293.JPEG
│  │   ├── ILSVRC2012_val_00002138.JPEG
│  │   ├── ......
│  ├── ......ValidationReplace DATA_DIR with your imagenet validation set path and MODEL_DIR with the checkpoint pathCUDA_VISIBLE_DEVICES=0 bash eval.sh /path/to/imagenet/val /path/to/checkpointLabel dataWe provide NFNet-F6 generated dense label map inGoogle DriveandBaiDu Yun(password: y6j2). As NFNet-F6 are based on pure ImageNet data, no extra training data is involved.TrainingTrain the LV-ViT-S:If only 4 GPUs are available,CUDA_VISIBLE_DEVICES=0,1,2,3 ./distributed_train.sh 4 /path/to/imagenet --model lvvit_s -b 256 --apex-amp --img-size 224 --drop-path 0.1 --token-label --token-label-data /path/to/label_data --token-label-size 14 --model-emaIf 8 GPUs are available:CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 /path/to/imagenet --model lvvit_s -b 128 --apex-amp --img-size 224 --drop-path 0.1 --token-label --token-label-data /path/to/label_data --token-label-size 14 --model-emaTrain the LV-ViT-M and LV-ViT-L (run on 8 GPUs):CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 /path/to/imagenet --model lvvit_m -b 128 --apex-amp --img-size 224 --drop-path 0.2 --token-label --token-label-data /path/to/label_data --token-label-size 14 --model-emaCUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 /path/to/imagenet --model lvvit_l -b 128 --lr 1.e-3 --aa rand-n3-m9-mstd0.5-inc1 --apex-amp --img-size 224 --drop-path 0.3 --token-label --token-label-data /path/to/label_data --token-label-size 14 --model-emaIf you want to train our LV-ViT on images with 384x384 resolution, please use--img-size 384 --token-label-size 24.Fine-tuningTo Fine-tune the pre-trained LV-ViT-S on images with 384x384 resolution:CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 /path/to/imagenet --model lvvit_s -b 64 --apex-amp --img-size 384 --drop-path 0.1 --token-label --token-label-data /path/to/label_data --token-label-size 24 --lr 5.e-6 --min-lr 5.e-6 --weight-decay 1.e-8 --finetune /path/to/checkpointTo Fine-tune the pre-trained LV-ViT-S on other datasets without token labeling:CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 ./distributed_train.sh 8 /path/to/dataset --model lvvit_s -b 64 --apex-amp --img-size 224 --drop-path 0.1 --token-label --token-label-size 14 --dense-weight 0.0 --num-classes $NUM_CLASSES --finetune /path/to/checkpointSegmentationOur Segmentation model are fully based upon theMMSegmentationToolkit. The model and config files are underseg/folder which follow the same folder structure. You can simply drop in these file to get start.gitclonehttps://github.com/open-mmlab/mmsegmentation# and installcpseg/mmseg/models/backbones/vit.pymmsegmentation/mmseg/models/backbones/
cp-rseg/configs/lvvitmmsegmentation/configs/# test upernet+lvvit_s (add --aug-test to test on multi scale)cdmmsegmentation
./tools/dist_test.shconfigs/lvvit/upernet_lvvit_s_512x512_160k_ade20k.py/path/to/checkpoint8--evalmIoU[--aug-test]BackboneMethodCrop sizeLr SchdmIoUmIoU(ms)Pixel Acc.ParamDownloadLV-ViT-SUperNet512x512160k47.948.683.144MlinkLV-ViT-MUperNet512x512160k49.450.683.577MlinkLV-ViT-LUperNet512x512160k50.951.884.1209MlinkVisualizationWe apply the visualization method in thisrepoto visualize the parts of the image that led to a certain classification for DeiT-Base and our LV-ViT-S. The parts of the image that used by the network to make the decision are highlighted in red.Label generationTo generate token label data for training:python3generate_label.py/path/to/imagenet/train/path/to/save/label_top5_train_nfnet--modeldm_nfnet_f6--pretrained--img-size576-b32--crop-pct1.0ReferenceIf you use this repo or find it useful, please consider citing:@article{jiang2021all,
  title={All Tokens Matter: Token Labeling for Training Better Vision Transformers},
  author={Jiang, Zihang and Hou, Qibin and Yuan, Li and Zhou, Daquan and Shi, Yujun and Jin, Xiaojie and Wang, Anran and Feng, Jiashi},
  journal={arXiv preprint arXiv:2104.10858},
  year={2021}
}Related projectsT2T-ViT,Re-labeling ImageNet,MMSegmentation,Transformer Explainability.