Transformer Knowledge Distillation for Efficient Semantic Segmentation [arxiv]
We propose the structural framework, TransKD, to distill the knowledge from feature maps and patch embeddings of vision transformers. TransKD enables the non-pretrained vision transformers perform on-par with the pretrained ones.
Environment: create a conda environment and activate it
conda create -n TransKD python=3.6
conda activate TransKD
Additional python pachages: poly scheduler and
pytorch == 1.7.1+cu92
torchvision == 0.8.2+cu92
mmsegmentation == 0.15.0
mmcv-full == 1.3.10
numpy
visdom
Datasets:
- Cityscapes: download
gtFine_trainvaltest.zipandleftImg8bit_trainvaltest.zipfrom cityscapes official website, then prepare the 19-class label with thecreateTrainIdLabelImgs.pyfrom cityscapesscripts. - ACDC： download
gt_trainval.zipandrgb_anon_trainvaltest.zipfrom ACDC official website.
Cityscapes:
| Network | #Params(M) | GFLOPs | mIoU(%) | weight |
|---|---|---|---|---|
| Teacher(B2) | 27.36 | 113.84 | 76.49 | Google Drive |
| Student(B0) | 3.72 | 13.67 | 55.86 | Google Drive |
| +TransKD-Base | 4.56 | 16.47 | 68.58 | Google Drive |
| +TransKD-GL | 5.22 | 16.80 | 68.87 | Google Drive |
| +TransKD-EA | 5.53 | 17.84 | 68.98 | Google Drive |
ACDC:
| Network | mIoU(%) | weight |
|---|---|---|
| Teacher(B2) | 69.34 | Google Drive |
| Student(B0) | 46.26 | Google Drive |
| +TransKD-Base | 58.56 | Google Drive |
| +TransKD-GL | 58.13 | Google Drive |
| +TransKD-EA | 59.09 | Google Drive |
Download pretrained weights (SegFormer and PVTv2) to the folder \train\ckpt_pretained\.
cd train
CUDA_VISIBLE_DEVICES=0 python TransKDBase.py --dataset cityscapes --dataset /path/to/cityscapes #--dataset ACDC --dataset /path/to/ACDC
CUDA_VISIBLE_DEVICES=0 python TransKD_GLMixer.py --dataset cityscapes --dataset /path/to/cityscapes #--dataset ACDC --dataset /path/to/ACDC
CUDA_VISIBLE_DEVICES=0 python TransKD_EA.py --dataset cityscapes --dataset /path/to/cityscapes #--dataset ACDC --dataset /path/to/ACDC
Download trained weights (Google Drive) to the folder \outputs.
cd eval
CUDA_VISIBLE_DEVICES=0 python eval_cityscapes_iou.py --distillation-type TransKDBase
CUDA_VISIBLE_DEVICES=0 python eval_cityscapes_color.py --distillation-type TransKDBase
CUDA_VISIBLE_DEVICES=0 python eval_ACDC_iou.py --distillation-type TransKDBase
CUDA_VISIBLE_DEVICES=0 python eval_ACDC_color.py --distillation-type TransKDBase
# distillation-type can be choices=['teacher','student','TransKDBase','TransKD_GL','TransKD_EA']
If you find this repo useful, please consider referencing the following paper [PDF]:
@article{liu2022transkd,
title={TransKD: Transformer Knowledge Distillation for Efficient Semantic Segmentation},
author={Liu, Ruiping and Yang, Kailun and Roitberg, Alina and Zhang, Jiaming and Peng, Kunyu and Liu, Huayao and Stiefelhagen, Rainer},
journal={arXiv preprint arXiv:2202.13393},
year={2022}
}
React
Typescript
Django
Laravel
Facebook
Microsoft
Google
Alibaba
D3
Tencent