Prune and Merge for Segmenter

The implementation of Prune and Merge Method for Segmenter.

Installation

The installation can follow the original segmenter link

Model Zoo

We release models of compressed Seg-L-Mask/16 with different compression rates.

ADE20K

Name	mIoU (SS)	Gflops	Resolution	FPS	Download
Seg-L-Mask/16	51.8	658	640x640	4.18	model	config
Seg-L-R08	51.9	528	640x640	5.25	model	config
Seg-L-R07	51.7	456	640x640	5.74	model	config
Seg-L-R05	50.2	364	640x640	6.85	model	config

Inference

Download one checkpoint with its configuration in a common folder, for example, seg_large_mask.

To evaluate on ADE20K, run the command:

# single-scale evaluation:
python -m segm.eval.miou seg_large_mask/checkpoint.pth ade20k --singlescale
# multi-scale evaluation:
python -m segm.eval.miou seg_large_mask/checkpoint.pth ade20k --multiscale

Compress

Compress Seg-Large-Mask/16 on ADE20K on 2 3090 GPU: (You should change the batch size in config.yml according to your GPU number)

python -m torch.distributed.launch --master_port=3001 --nproc_per_node=2 --use_env \
prune_finetune.py --prune --pretrain-dir path/to/pretrain/dir \
--dataset ade20k --log-dir path/to/log/dir \
--backbone vit_large_patch16_384 --decoder mask_transformer --iter-num 100

Finetune

Finetune Seg-Large-Mask/16 on ADE20K on 8 3090 GPU:

python -m torch.distributed.launch --master_port=3001 --nproc_per_node=8 --use_env \
prune_finetune.py --eval --pretrain-dir path/to/pretrain/dir \
--dataset ade20k --log-dir path/to/log/dir \
--distill-type soft --teacher-path path/to/pretrain/model \
--backbone vit_large_patch16_384 --decoder mask_transformer --iter-num 0  \
--scheduler cosine -lr 1e-4 --alpha 1.0 --tau 20 --weight-decay 0.0001 

Acknowledgements

Our implementation is base on segmenter