Knowledge Distillation from A Stronger Teacher (DIST)

Official implementation of paper "Knowledge Distillation from A Stronger Teacher" (DIST), NeurIPS 2022.
By Tao Huang, Shan You, Fei Wang, Chen Qian, Chang Xu.

🔥 DIST: a simple and effective KD method.

Updates

• December 27, 2022: Update CIFAR-100 distillation code and logs.

• September 20, 2022: Release code for semantic segmentation task.

• September 15, 2022: DIST was accepted by NeurIPS 2022!

• May 30, 2022: Code for object detection is available.

• May 27, 2022: Code for ImageNet classification is available.

Getting started

Clone training code

git clone https://github.com/hunto/DIST_KD.git --recurse-submodules
cd DIST_KD

The loss function of DIST is in classification/lib/models/losses/dist_kd.py.

• classification: prepare your environment and datasets following the README.md in classification.
• object detection: coming soon.
• semantic segmentation: coming soon.

Reproducing our results

ImageNet

cd classification
sh tools/dist_train.sh 8 ${CONFIG} ${MODEL} --teacher-model ${T_MODEL} --experiment ${EXP_NAME}

• Baseline settings (R34-R101 and R50-MBV1):

  CONFIG=configs/strategies/distill/resnet_dist.yaml
  

  Student	Teacher	DIST	MODEL	T_MODEL	Log	Ckpt
  ResNet-18 (69.76)	ResNet-34 (73.31)	72.07	tv_resnet18	tv_resnet34	log	ckpt
  MobileNet V1 (70.13)	ResNet-50 (76.16)	73.24	mobilenet_v1	tv_resnet50	log	ckpt

• Stronger teachers (R18 and R34 students with various ResNet teachers):

  Student	Teacher	KD (T=4)	DIST
  ResNet-18 (69.76)	ResNet-34 (73.31)	71.21	72.07
  ResNet-18 (69.76)	ResNet-50 (76.13)	71.35	72.12
  ResNet-18 (69.76)	ResNet-101 (77.37)	71.09	72.08
  ResNet-18 (69.76)	ResNet-152 (78.31)	71.12	72.24
  ResNet-34 (73.31)	ResNet-50 (76.13)	74.73	75.06
  ResNet-34 (73.31)	ResNet-101 (77.37)	74.89	75.36
  ResNet-34 (73.31)	ResNet-152 (78.31)	74.87	75.42

• Stronger training strategies:

  CONFIG=configs/strategies/distill/dist_b2.yaml
  

  ResNet-50-SB: stronger ResNet-50 trained by TIMM (ResNet strikes back) .

  Student	Teacher	KD (T=4)	DIST	MODEL	T_MODEL	Log
  ResNet-18 (73.4)	ResNet-50-SB (80.1)	72.6	74.5	tv_resnet18	timm_resnet50	log
  ResNet-34 (76.8)	ResNet-50-SB (80.1)	77.2	77.8	tv_resnet34	timm_resnet50	log
  MobileNet V2 (73.6)	ResNet-50-SB (80.1)	71.7	74.4	tv_mobilenet_v2	timm_resnet50	log
  EfficientNet-B0 (78.0)	ResNet-50-SB (80.1)	77.4	78.6	timm_tf_efficientnet_b0	timm_resnet50	log
  ResNet-50 (78.5)	Swin-L (86.3)	80.0	80.2	tv_resnet50	timm_swin_large_patch4_window7_224	log ckpt
  Swin-T (81.3)	Swin-L (86.3)	81.5	82.3	-	-	log

  • Swin-L student: We implement our DIST on the official code of Swin-Transformer.

CIFAR-100

Download and extract the teacher checkpoints to your disk, then specify the path of the corresponding checkpoint pth file using --teacher-ckpt:

cd classification
sh tools/dist_train.sh 1 configs/strategies/distill/dist_cifar.yaml ${MODEL} --teacher-model ${T_MODEL} --experiment ${EXP_NAME} --teacher-ckpt ${CKPT}

NOTE: For MobileNetV2, ShuffleNetV1, and ShuffleNetV2, lr and warmup-lr should be 0.01:

sh tools/dist_train.sh 1 configs/strategies/distill/dist_cifar.yaml ${MODEL} --teacher-model ${T_MODEL} --experiment ${EXP_NAME} --teacher-ckpt ${CKPT} --lr 0.01 --warmup-lr 0.01

COCO Detection

The training code is in MasKD/mmrazor. An example to train cascade_mask_rcnn_x101-fpn_r50:

sh tools/mmdet/dist_train_mmdet.sh configs/distill/dist/dist_cascade_mask_rcnn_x101-fpn_x50_coco.py 8 work_dirs/dist_cmr_x101-fpn_x50

Cityscapes Segmentation

Detailed instructions of reproducing our results are in segmentation folder (README).

License

This project is released under the Apache 2.0 license.

Citation

@article{huang2022knowledge,
  title={Knowledge Distillation from A Stronger Teacher},
  author={Huang, Tao and You, Shan and Wang, Fei and Qian, Chen and Xu, Chang},
  journal={arXiv preprint arXiv:2205.10536},
  year={2022}
}