RepDistiller

本仓库：

(1) 包含以下 ICLR 2020 论文的实现：

"Contrastive Representation Distillation"（对比表示蒸馏，CRD）。论文，项目主页。

(2) 在 PyTorch 中基准测试了 12 种最先进的知识蒸馏（Knowledge Distillation，KD）方法，包括：

(KD) - Distilling the Knowledge in a Neural Network
(FitNet) - Fitnets: hints for thin deep nets
(AT) - Paying More Attention to Attention: Improving the Performance of Convolutional Neural Networks via Attention Transfer
(SP) - Similarity-Preserving Knowledge Distillation
(CC) - Correlation Congruence for Knowledge Distillation
(VID) - Variational Information Distillation for Knowledge Transfer
(RKD) - Relational Knowledge Distillation
(PKT) - Probabilistic Knowledge Transfer for deep representation learning
(AB) - Knowledge Transfer via Distillation of Activation Boundaries Formed by Hidden Neurons
(FT) - Paraphrasing Complex Network: Network Compression via Factor Transfer
(FSP) - A Gift from Knowledge Distillation: Fast Optimization, Network Minimization and Transfer Learning
(NST) - Like what you like: knowledge distill via neuron selectivity transfer

安装

本仓库在 Ubuntu 16.04.5 LTS、Python 3.5、PyTorch 0.4.0 和 CUDA 9.0 环境下测试通过。但应该可以在 PyTorch >=0.4.0 的较新版本上运行。

运行

1. 通过以下命令获取预训练的教师模型：

   sh scripts/fetch_pretrained_teachers.sh
   

   该命令将下载模型并保存到 save/models 目录

2. 按照 scripts/run_cifar_distill.sh 中的命令运行蒸馏。运行 Geoffrey 原始知识蒸馏（KD）的示例如下：

   python train_student.py --path_t ./save/models/resnet32x4_vanilla/ckpt_epoch_240.pth --distill kd --model_s resnet8x4 -r 0.1 -a 0.9 -b 0 --trial 1
   

   各参数说明如下：

   • --path_t：指定教师模型的路径
   • --model_s：指定学生模型，查看 models/__init__.py 了解可用的模型类型
   • --distill：指定蒸馏方法
   • -r：logit 与真实标签之间交叉熵损失的权重，默认值：1
   • -a：KD 损失的权重，默认值：None
   • -b：其他蒸馏损失的权重，默认值：None
   • --trial：指定实验 ID，用于区分多次运行

   因此，运行 CRD 的命令如下：

   python train_student.py --path_t ./save/models/resnet32x4_vanilla/ckpt_epoch_240.pth --distill crd --model_s resnet8x4 -a 0 -b 0.8 --trial 1
   

3. 将蒸馏目标与 KD 结合使用，只需将 -a 设为非零值，示例如下（CRD 与 KD 结合）：

   python train_student.py --path_t ./save/models/resnet32x4_vanilla/ckpt_epoch_240.pth --distill crd --model_s resnet8x4 -a 1 -b 0.8 --trial 1     
   

4. （可选）从头开始训练教师网络。示例命令见 scripts/run_cifar_vanilla.sh

注意：默认设置为单 GPU 训练。如需使用多 GPU 运行本仓库，可能需要调整学习率，根据经验，学习率需要随批量大小（batch size）线性放大，详见此论文。

CIFAR-100 基准测试结果：

性能以分类准确率（%）衡量

1. 教师和学生模型为相同架构类型。

教师 学生	wrn-40-2 wrn-16-2	wrn-40-2 wrn-40-1	resnet56 resnet20	resnet110 resnet20	resnet110 resnet32	resnet32x4 resnet8x4	vgg13 vgg8
教师 学生	75.61 73.26	75.61 71.98	72.34 69.06	74.31 69.06	74.31 71.14	79.42 72.50	74.64 70.36
KD	74.92	73.54	70.66	70.67	73.08	73.33	72.98
FitNet	73.58	72.24	69.21	68.99	71.06	73.50	71.02
AT	74.08	72.77	70.55	70.22	72.31	73.44	71.43
SP	73.83	72.43	69.67	70.04	72.69	72.94	72.68
CC	73.56	72.21	69.63	69.48	71.48	72.97	70.71
VID	74.11	73.30	70.38	70.16	72.61	73.09	71.23
RKD	73.35	72.22	69.61	69.25	71.82	71.90	71.48
PKT	74.54	73.45	70.34	70.25	72.61	73.64	72.88
AB	72.50	72.38	69.47	69.53	70.98	73.17	70.94
FT	73.25	71.59	69.84	70.22	72.37	72.86	70.58
FSP	72.91	N/A	69.95	70.11	71.89	72.62	70.23
NST	73.68	72.24	69.60	69.53	71.96	73.30	71.53
CRD	75.48	74.14	71.16	71.46	73.48	75.51	73.94

2. 教师和学生模型为不同架构类型。

教师 学生	vgg13 MobileNetV2	ResNet50 MobileNetV2	ResNet50 vgg8	resnet32x4 ShuffleNetV1	resnet32x4 ShuffleNetV2	wrn-40-2 ShuffleNetV1
教师 学生	74.64 64.60	79.34 64.60	79.34 70.36	79.42 70.50	79.42 71.82	75.61 70.50
KD	67.37	67.35	73.81	74.07	74.45	74.83
FitNet	64.14	63.16	70.69	73.59	73.54	73.73
AT	59.40	58.58	71.84	71.73	72.73	73.32
SP	66.30	68.08	73.34	73.48	74.56	74.52
CC	64.86	65.43	70.25	71.14	71.29	71.38
VID	65.56	67.57	70.30	73.38	73.40	73.61
RKD	64.52	64.43	71.50	72.28	73.21	72.21
PKT	67.13	66.52	73.01	74.10	74.69	73.89
AB	66.06	67.20	70.65	73.55	74.31	73.34
FT	61.78	60.99	70.29	71.75	72.50	72.03
NST	58.16	64.96	71.28	74.12	74.68	74.89
CRD	69.73	69.11	74.30	75.11	75.65	76.05

引用

如果本仓库对您的研究有所帮助，请考虑引用该论文：

@inproceedings{tian2019crd,
  title={Contrastive Representation Distillation},
  author={Yonglong Tian and Dilip Krishnan and Phillip Isola},
  booktitle={International Conference on Learning Representations},
  year={2020}
}

如有任何问题，请联系 Yonglong Tian（yonglong@mit.edu）。

致谢

感谢 Baoyun Peng 提供 CC（Correlation Congruence，相关性一致性）的代码，以及 Frederick Tung 帮助我们验证 SP（Similarity Preserving，相似性保持）的重新实现。同时感谢其他论文的作者公开他们的代码。

RepDistiller 快速上手指南

RepDistiller 是知识蒸馏领域的综合基准框架，实现了 ICLR 2020 论文 "Contrastive Representation Distillation" (CRD)，并集成了 12 种主流知识蒸馏方法。

---

环境准备

项目	要求
操作系统	Ubuntu 16.04+（推荐）
Python	3.5+
PyTorch	≥ 0.4.0
CUDA	9.0+

注：框架兼容更高版本 PyTorch，但多 GPU 训练时需按线性比例调整学习率。

---

安装步骤

# 1. 克隆仓库
git clone https://github.com/HobbitLong/RepDistiller.git
cd RepDistiller

# 2. 安装依赖（建议先创建虚拟环境）
pip install torch torchvision

# 3. 下载预训练教师模型
sh scripts/fetch_pretrained_teachers.sh

模型将自动保存至 save/models/ 目录。

---

基本使用

单方法蒸馏（以 CRD 为例）

python train_student.py \
    --path_t ./save/models/resnet32x4_vanilla/ckpt_epoch_240.pth \
    --distill crd \
    --model_s resnet8x4 \
    -a 0 \
    -b 0.8 \
    --trial 1

组合蒸馏（CRD + KD）

python train_student.py \
    --path_t ./save/models/resnet32x4_vanilla/ckpt_epoch_240.pth \
    --distill crd \
    --model_s resnet8x4 \
    -a 1 \
    -b 0.8 \
    --trial 1

经典 KD 方法

python train_student.py \
    --path_t ./save/models/resnet32x4_vanilla/ckpt_epoch_240.pth \
    --distill kd \
    --model_s resnet8x4 \
    -r 0.1 \
    -a 0.9 \
    -b 0 \
    --trial 1

---

参数说明

参数	说明	默认值
--path_t	教师模型路径	必填
--model_s	学生模型类型	必填
--distill	蒸馏方法（kd/crd/at/sp 等）	必填
-r	交叉熵损失权重	1
-a	KD 损失权重	None
-b	其他蒸馏损失权重	None
--trial	实验编号（区分多次运行）	0

可用模型类型见 models/__init__.py，完整命令参考 scripts/run_cifar_distill.sh。

---

可选：从头训练教师网络

# 参考脚本中的命令
sh scripts/run_cifar_vanilla.sh