Vision GNN

By Kai Han, Yunhe Wang, Jianyuan Guo, Yehui Tang and Enhua Wu. NeurIPS 2022. [arXiv link]

Requirements

Pytorch 1.7.0, timm 0.3.2, torchprofile 0.0.4, apex

Pretrained models

Model	Params (M)	FLOPs (B)	Top-1	Github Release
ViG-Ti	7.1	1.3	73.9	Github Release
ViG-S	22.7	4.5	80.4	Github Release
ViG-B	86.8	17.7	82.3	Github Release

Pyramid ViG

Model	Params (M)	FLOPs (B)	Top-1	BaiduDisk URL	Github Release
Pyramid ViG-Ti	10.7	1.7	78.5	[BaiduDisk], Password: chae	Github Release
Pyramid ViG-S	27.3	4.6	82.1	[BaiduDisk], Password: 81mg	Github Release
Pyramid ViG-M	51.7	8.9	83.1	[BaiduDisk], Password: prd3	Github Release
Pyramid ViG-B	82.6	16.8	83.7	[BaiduDisk], Password: rgm4	Github Release

Evaluation

Data preparation follows the official pytorch example

Evaluate example:

python train.py /path/to/imagenet/ --model pvig_s_224_gelu -b 128 --pretrain_path /path/to/pretrained/model/ --evaluate

Training

Training ViG on 8 GPUs:

python -m torch.distributed.launch --nproc_per_node=8 train.py /path/to/imagenet/ --model vig_s_224_gelu --sched cosine --epochs 300 --opt adamw -j 8 --warmup-lr 1e-6 --mixup .8 --cutmix 1.0 --model-ema --model-ema-decay 0.99996 --aa rand-m9-mstd0.5-inc1 --color-jitter 0.4 --warmup-epochs 20 --opt-eps 1e-8 --repeated-aug --remode pixel --reprob 0.25 --amp --lr 2e-3 --weight-decay .05 --drop 0 --drop-path .1 -b 128 --output /path/to/save/models/

Training Pyramid ViG on 8 GPUs:

python -m torch.distributed.launch --nproc_per_node=8 train.py /path/to/imagenet/ --model pvig_s_224_gelu --sched cosine --epochs 300 --opt adamw -j 8 --warmup-lr 1e-6 --mixup .8 --cutmix 1.0 --model-ema --model-ema-decay 0.99996 --aa rand-m9-mstd0.5-inc1 --color-jitter 0.4 --warmup-epochs 20 --opt-eps 1e-8 --repeated-aug --remode pixel --reprob 0.25 --amp --lr 2e-3 --weight-decay .05 --drop 0 --drop-path .1 -b 128 --output /path/to/save/models/

Citation

@inproceedings{han2022vig,
  title={Vision GNN: An Image is Worth Graph of Nodes}, 
  author={Kai Han and Yunhe Wang and Jianyuan Guo and Yehui Tang and Enhua Wu},
  booktitle={NeurIPS},
  year={2022}
}

Acknowledgement

This repo partially uses code from deep_gcns_torch and timm.

메모

srun -J vig -p A5000 -N 1 -n 1 --time 01:00:00 --gres=gpu:8 --cpus-per-task=10 --pty /bin/bash -l

conda activate vig

cd train_vig

torchrun --nproc_per_node=8 ./train.py --world-size 8 --workers 10 --batch-size 32

#SBATCH -p 4A100

#SBATCH -q 4A100

euna-jeong20/project