Update

Pipeline

Requirements

torch>=1.12.0torchvision>=0.13.0timm==0.9.2

How to apply ToE

 model = create_model(     args.model,     pretrained=False,     num_classes=args.nb_classes,     drop_rate=args.drop,     drop_path_rate=args.drop_path,     drop_block_rate=None,     img_size=args.input_size, ) # ---------------------# model.token_select = TokenSelect(     expansion_step=args.expansion_step,     keep_rate=args.keep_rate,     initialization_keep_rate=args.initialization_keep_rate,     expansion_multiple_stage=args.expansion_multiple_stage,     distance=args.distance, ) # ---------------------#

Data prepare

ImageNet-1k

│ILSVRC2012/├──train/│  ├── n01440764│  │   ├── n01440764_10026.JPEG│  │   ├── n01440764_10027.JPEG│  │   ├── ......│  ├── ......├──val/│  ├── n01440764│  │   ├── ILSVRC2012_val_00000293.JPEG│  │   ├── ILSVRC2012_val_00002138.JPEG│  │   ├── ......│  ├── ......

COCO

path/to/coco/  annotations/  # annotation json files  train2017/    # train images  val2017/      # val images

Label data for LV-ViT

Pretrained model prepare

Pretrained DeiT model for YOLOS

Training

DeiT-small

cd ToE/deitbash run/imagenet_small_run1.sh

result_dir=[your_result_path]dataset_dir=[your_dataset_path]device=0,1,2,3master_port=6666CUDA_VISIBLE_DEVICES=$device torchrun --nproc_per_node=4 --master_port=$master_port main.py \--patch-size 16 \--model deit_small_patch16_224 \--batch-size 256 \--data-path $dataset_dir \--output_dir $result_dir \--num_workers 8 \--seed 3407 \--expansion-step 0 100 200 \--keep-rate 0.5 0.75 1.0 \--initialization-keep-rate 0.25 \--expansion-multiple-stage 2 \

LV-ViT-S

cd ToE/lvvitbash run/imagenet_small_run1.sh

#!/bin/bashresult_dir=[your_result_path]dataset_dir=[your_dataset_path]label_dir=[your_token_labeling_dataset_path]device=0,1,2,3master_port=6666shiftCUDA_VISIBLE_DEVICES=$device python3 -m torch.distributed.launch --nproc_per_node=4 main.py "$@" $dataset_dir \--output $result_dir \--model lvvit_s \-b 256 \--img-size 224 \--drop-path 0.1 \--token-label \--token-label-data $label_dir \--token-label-size 14 \--model-ema \--apex-amp \--expansion-step 0 100 200 \--keep-rate 0.4 0.7 1.0 \--initialization-keep-rate 0.2 \--expansion-multiple-stage 2 \

EfficientTrain (DeiT-small)

cd ToE/EfficientTrainbash run/imagenet_small_run1.sh

result_dir=[your_result_path]dataset_dir=[your_dataset_path]CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python ET_training.py \--data_path $dataset_dir \--output_dir $result_dir \--model deit_small_patch16_224 \--final_bs 256 --epochs 300 \--num_gpus 8 --num_workers 8 \

"deit_tiny_patch16_224": " --use_amp true --clip_grad 5.0 \--expansion-step 0 100 200 --keep-rate 0.6 0.8 1.0 \--initialization-keep-rate 0.3 --expansion-multiple-stage 2 "

YOLOS-S

cd ToE/YOLOSbash run/coco_small_run1.sh

result_dir=[your_result_path]dataset_dir=[your_dataset_path]pretrain_path=[your_pretrain_model_path]device=0,1,2,3,4,5,6,7master_port=6666CUDA_VISIBLE_DEVICES=$device python -m torch.distributed.launch --nproc_per_node=8 --master_port=$master_port --use_env main.py \    --coco_path $dataset_dir \    --batch_size 1 \    --lr 2.5e-5 \    --epochs 150 \    --backbone_name small \    --pre_trained $pretrain_path \    --eval_size 800 \    --init_pe_size 512 864 \    --mid_pe_size 512 864 \    --output_dir $result_dir \    --num_workers 8 \    --expansion-step 5 50 100 \    --keep-rate 0.5 0.75 1.0 \    --initialization-keep-rate 0.25 \    --expansion-multiple-stage 2 \

Acknowledgements

Citation

@article{huang2024general,  title={A General and Efficient Training for Transformer via Token Expansion},  author={Huang, Wenxuan and Shen, Yunhang and Xie, Jiao and Zhang, Baochang and He, Gaoqi and Li, Ke and Sun, Xing and Lin, Shaohui},  journal={CVPR},  year={2024}}