trellixvulnteam/unixkd_0y78

UNIXKD

This repo is the implementation of CIFAR100 part in paper Computation-Efficient Knowledge Distillation via Uncertainty-Aware Mixup.

Prerequisite

This repo is tested with Ubuntu 16.04.5, Python 3.7, PyTorch 1.5.0, CUDA 10.2. Make sure to install pytorch, torchvision, tensorboardX, numpy before using this repo.

Running

Teacher Training

An example of teacher training is:

python teacher.py --arch wrn_40_2 --lr 0.05 --gpu-id 0

where you can specify the architecture via flag --arch

You can also download all the pre-trained teacher models here. If you want to run student.py directly, you have to re-organise the directory. For instance, when you download vgg13.pth, you have to make a directory for it, say teacher_vgg13, and then make a new directory ckpt inside teacher_vgg13. Move the vgg13.pth into teacher_vgg13/ckpt and rename it as best.pth. If you want a simpler way to use pre-trained model, you can edit the code in student_v0.py (line 96).

Student Training

An example of student training is:

python student_v0.py --teacher-path ./experiments/teacher_wrn_40_2/ --student-arch wrn_16_2 --lr 0.05 --gpu-id 0

The meanings of flags are:

--teacher-path: teacher's checkpoint path. Automatically search the checkpoint containing 'best' keyword in its name.

--student-arch: student's architecture.

All the commands can be found in command/command.sh

Results (Top-1 Acc) on CIFAR100

The accuracies are slightly different from those in the paper. Average on 4 runs.

Cross-Architecture

Teacher Student	vgg13 MobieleNetV2	ResNet50 MobileNetV2	ResNet50 vgg8	resnet32x4 ShuffleV1	resnet32x4 ShuffleV2	wrn40-2 ShuffleV1
Teacher Student	75.38 65.79	79.10 65.79	79.10 70.68	79.63 70.77	79.63 73.12	76.46 70.77
Ft/Fs	38.17	174.00	13.56	27.16	23.49	8.22
KD Computation	67.94 100%	68.33 100%	73.43 100%	74.52 100%	75.07 100%	76.04 100%
UNIXKD Computation	68.09 77.49%	68.76 75.57%	74.02 81.43%	76.48 78.43%	76.86 78.92%	77.06 84.79%

Similar-Architecture

Teacher Student	wrn40-2 wrn16-2	wrn40-2 wrn40-1	resnet56 resnet20	resnet32x4 resnet8x4	vgg13 vgg8
Teacher Student	76.46 73.64	76.46 72.24	73.44 69.63	79.63 72.51	75.38 70.68
Ft/Fs	3.25	3.93	3.06	6.12	2.97
KD Computation	75.40 100%	73.77 100%	70.72 100%	73.34 100%	73.38 100%
UNIXKD Computation	75.40 94.06%	74.38 91.88%	70.41 94.76%	74.86 87.32%	73.55 95.10%

Citation

If you find this repo useful for your research, please consider citing the paper

@misc{xu2020unixkd,
      title={Computation-Efficient Knowledge Distillation via Uncertainty-Aware Mixup}, 
      author={Guodong Xu and Ziwei Liu and Chen Change Loy},
      year={2020},
      eprint={2012.09413},
      archivePrefix={arXiv},
      primaryClass={cs.CV}
}