Official implementation of paper "Masked Distillation with Receptive Tokens" (MasKD), ICLR 2023.

By Tao Huang*, Yuan Zhang*, Shan You, Fei Wang, Chen Qian, Jian Cao, Chang Xu.

🔥 MasKD: better and more general feature distillation method for dense prediction tasks (e.g., detection and segmentation).

Configs for KD are available in mmrazor folders and student checkpoints are provided.

Code for mask learning and KD is available in mmdetection and mmrazor folders.

We provide the learned pretrained mask tokens in our experiments at release.

This repo uses MMRazor as the knowledge distillation toolkit. For environment setup, please see mmrazor/README.md.

Train student:

cd mmrazor
sh tools/mmdet/dist_train_mmdet.sh ${CONFIG} 8 ${WORK_DIR}

Example for reproducing our cascade_mask_rcnn_x101-fpn_r50 result:

sh tools/mmdet/dist_train_mmdet.sh configs/distill/maskd/cascade_mask_rcnn_x101-fpn_x50_coco.py 8 work_dirs/maskd_cmr_x101-fpn_x50

• Baseline settings:

  Student	Teacher	MasKD	Config	Log	CheckPoint
  Faster RCNN-R50 (38.4)	Faster RCNN-R101 (39.8)	41.0	config	log	GoogleDrive
  RetinaNet-R50 (37.4)	RetinaNet-R101 (38.9)	39.9	config	log	GoogleDrive
  FCOS-R50 (38.5)	FCOS-R101 (40.8)	42.9	config	log	GoogleDrive

• Stronger teachers:

  Student	Teacher	MasKD	Config	Log	CheckPoint
  Faster RCNN-R50 (38.4)	Cascade Mask RCNN-X101 (45.6)	42.7	config	log	GoogleDrive
  RetinaNet-R50 (37.4)	RetinaNet-X101 (41.0)	41.0	config	log	GoogleDrive
  RepPoints-R50 (38.6)	RepPoints-R101 (44.2)	42.5	config	log	GoogleDrive

You can train your own mask tokens with the code provided in mmdetection folder. Please check mmdetection/README.md for detailed instructions.

For semantic segmentation, please see segmentation folder.

This project is released under the Apache 2.0 license.

@inproceedings{
huang2023masked,
title={Masked Distillation with Receptive Tokens},
author={Tao Huang and Yuan Zhang and Shan You and Fei Wang and Chen Qian and Jian Cao and Chang Xu},
booktitle={The Eleventh International Conference on Learning Representations },
year={2023},
url={https://openreview.net/forum?id=mWRngkvIki3}
}