- Train a PVT model (e.g., PVTv2 B0) with standard self-attention under 100 epochs. The model is initialized with corresponding pre-trained models in PVT.
# train with 8 GPUs
python -m torch.distributed.launch --nproc_per_node=8 --master_port=1236 \
--use_env main.py \
--config configs/pvt_v2/pvt_v2_b0_msa.py \
--batch-size 32 \
--data-path [path of imagenet] \
--data-set IMNET \
--epochs 100 \
--lr 5e-5 \
--warmup-lr 1e-7 \
--min-lr 1e-6 \
--finetune [path of pvt_v2 pre-trained models] \
--output_dir [output path of msa finetuned models] \- Convert MSA to linear attention and reparameterize all MatMuls with add layers
# train with 8 GPUs
python -m torch.distributed.launch --nproc_per_node=8 --master_port=1236 \
--use_env main.py \
--config configs/pvt_v2/pvt_v2_b0_LinAngular.py \
--batch-size 32 \
--data-path [path of imagenet] \
--data-set IMNET \
--epochs 100 \
--lr 5e-5 \
--warmup-lr 1e-7 \
--min-lr 1e-6 \
--finetune [path of msa finetuned models (from step 1)] \
--output_dir [output path of finetuned models with linear attention where all MatMuls are replaced with add layers] \- Reparameterize MLPs and linear projection layers in attention with MoE layers
# train with 8 GPUs
python -m torch.distributed.launch --nproc_per_node=8 --master_port=1236 \
--use_env main.py \
--config configs/pvt_v2/pvt_v2_b0_LinAngular.py \
--batch-size 32 \
--data-path [path of imagenet] \
--data-set IMNET \
--epochs 100 \
--lr 1e-5 \
--warmup-lr 5e-8 \
--min-lr 5e-7 \
--moe_attn \
--moe_mlp \
--finetune [path of finetuned models from step 2] \
--output_dir [output path of ShiftAddViT models] \To evaluate a model, you can
python -m torch.distributed.launch --nproc_per_node=1 --master_port=1236 \
--use_env main.py \
--config configs/pvt_v2/pvt_v2_b0_LinAngular.py \
--batch-size 32 \
--data-path [path of imagenet] \
--data-set IMNET \
--resume [path of finetuned ShiftAddViT modelss] \
--moe_attn \
--moe_mlp \
--evalTo speedup model with TVM kernel, you can
python main.py \
--config configs/pvt_v2/pvt_v2_b0_LinAngular.py \
--batch-size 32 \
--data-path [path of imagenet] \
--data-set IMNET \
--resume [path of finetuned ShiftAddViT modelss] \
--moe_attn \
--moe_mlp \
--tvm_tuneTo test latency on pytorch:
python main.py \
--config configs/pvt_v2/pvt_v2_b0_LinAngular.py \
--batch-size 32 \
--data-path [path of imagenet] \
--data-set IMNET \
--resume [path of finetuned ShiftAddViT modelss] \
--moe_attn \
--moe_mlp \
--throughputTo test latency with tuned TVM model (run TVM tune first):
python main.py \
--config configs/pvt_v2/pvt_v2_b0_LinAngular.py \
--batch-size 32 \
--data-path [path of imagenet] \
--data-set IMNET \
--resume [path of finetuned ShiftAddViT modelss] \
--moe_attn \
--moe_mlp \
--tvm_throughputTo analyse the energy cost of model:
python main.py \
--config configs/pvt_v2/pvt_v2_b0_LinAngular.py \
--batch-size 32 \
--data-path [path of imagenet] \
--data-set IMNET \
--resume [path of finetuned ShiftAddViT modelss] \
--moe_attn \
--moe_mlp \
--cal_energy