这个仓库是

(1) 一个 PyTorch 库，提供了在主流计算机视觉基准上的经典知识蒸馏算法，

(2) CVPR-2022 论文 解耦知识蒸馏 的官方实现。

(3) ICCV-2023 论文 DOT：面向蒸馏的训练器 的官方实现。

DOT：面向蒸馏的训练器

框架

主要基准测试结果

在 CIFAR-100 上：

教师 学生	ResNet32x4 ResNet8x4	VGG13 VGG8	ResNet32x4 ShuffleNet-V2
KD	73.33	72.98	74.45
KD+DOT	75.12	73.77	75.55

在 Tiny-ImageNet 上：

教师 学生	ResNet18 MobileNet-V2	ResNet18 ShuffleNet-V2
KD	58.35	62.26
KD+DOT	64.01	65.75

在 ImageNet 上：

教师 学生	ResNet34 ResNet18	ResNet50 MobileNet-V1
KD	71.03	70.50
KD+DOT	71.72	73.09

解耦知识蒸馏

框架与性能

主要基准测试结果

在 CIFAR-100 上：

教师 学生	ResNet56 ResNet20	ResNet110 ResNet32	ResNet32x4 ResNet8x4	WRN-40-2 WRN-16-2	WRN-40-2 WRN-40-1	VGG13 VGG8
KD	70.66	73.08	73.33	74.92	73.54	72.98
DKD	71.97	74.11	76.32	76.23	74.81	74.68

教师 学生	ResNet32x4 ShuffleNet-V1	WRN-40-2 ShuffleNet-V1	VGG13 MobileNet-V2	ResNet50 MobileNet-V2	ResNet32x4 MobileNet-V2
KD	74.07	74.83	67.37	67.35	74.45
DKD	76.45	76.70	69.71	70.35	77.07

在 ImageNet 上：

教师 学生	ResNet34 ResNet18	ResNet50 MobileNet-V1
KD	71.03	70.50
DKD	71.70	72.05

MDistiller

简介

MDistiller 支持在 CIFAR-100、ImageNet 和 MS-COCO 上的以下蒸馏方法：

方法	论文链接	CIFAR-100	ImageNet	MS-COCO
KD	https://arxiv.org/abs/1503.02531	✓	✓
FitNet	https://arxiv.org/abs/1412.6550	✓
AT	https://arxiv.org/abs/1612.03928	✓	✓
NST	https://arxiv.org/abs/1707.01219	✓
PKT	https://arxiv.org/abs/1803.10837	✓
KDSVD	https://arxiv.org/abs/1807.06819	✓
OFD	https://arxiv.org/abs/1904.01866	✓	✓
RKD	https://arxiv.org/abs/1904.05068	✓
VID	https://arxiv.org/abs/1904.05835	✓
SP	https://arxiv.org/abs/1907.09682	✓
CRD	https://arxiv.org/abs/1910.10699	✓	✓
ReviewKD	https://arxiv.org/abs/2104.09044	✓	✓	✓
DKD	https://arxiv.org/abs/2203.08679	✓	✓	✓

安装

环境要求：

• Python 3.6
• PyTorch 1.9.0
• torchvision 0.10.0

安装包：

sudo pip3 install -r requirements.txt
sudo python3 setup.py develop

使用指南

0. 使用 Wandb 作为日志记录器

• 注册地址： https://wandb.ai/home。
• 如果您不想使用 Wandb 作为日志记录器，请在 mdistiller/engine/cfg.py 中将 CFG.LOG.WANDB 设置为 False。

1. 评估

• 您可以评估我们提供的模型或您自己训练的模型的性能。

• 我们的模型位于 https://github.com/megvii-research/mdistiller/releases/tag/checkpoints,请将检查点下载到 ./download_ckpts。

• 如果要在 ImageNet 上测试模型，请从 https://image-net.org/ 下载数据集，并将其放入 ./data/imagenet。

  # 评估教师模型
  python3 tools/eval.py -m resnet32x4 # resnet32x4 在 cifar100 上
  python3 tools/eval.py -m ResNet34 -d imagenet # ResNet34 在 imagenet 上
  
  # 评估学生模型
  python3 tools/eval.p -m resnet8x4 -c download_ckpts/dkd_resnet8x4 # dkd-resnet8x4 在 cifar100 上
  python3 tools/eval.p -m MobileNetV1 -c download_ckpts/imgnet_dkd_mv1 -d imagenet # dkd-mv1 在 imagenet 上
  python3 tools/eval.p -m model_name -c output/your_exp/student_best # 您自己的检查点
  

2. 在 CIFAR-100 上训练

• 从 https://github.com/megvii-research/mdistiller/releases/tag/checkpoints 下载 cifar_teachers.tar，并使用 tar xvf cifar_teachers.tar 将其解压到 ./download_ckpts。

  # 以我们的 DKD 方法为例。
  python3 tools/train.py --cfg configs/cifar100/dkd/res32x4_res8x4.yaml
  
  # 您也可以在命令行中更改设置
  python3 tools/train.py --cfg configs/cifar100/dkd/res32x4_res8x4.yaml SOLVER.BATCH_SIZE 128 SOLVER.LR 0.1
  

3. 在 ImageNet 上训练

• 从 https://image-net.org/ 下载数据集，并将其放入 ./data/imagenet。

  # 以我们的 DKD 方法为例。
  python3 tools/train.py --cfg configs/imagenet/r34_r18/dkd.yaml
  

4. 在 MS-COCO 上训练

• 请参阅 detection.md

5. 扩展：可视化

• Jupyter 笔记本：tsne 和 correlation_matrices

自定义知识蒸馏方法

1. 在 mdistiller/distillers/ 目录下创建一个 Python 文件，并定义您的蒸馏器。

from ._base import Distiller

class MyDistiller(Distiller):
    def __init__(self, student, teacher, cfg):
        super(MyDistiller, self).__init__(student, teacher)
        self.hyper1 = cfg.MyDistiller.hyper1
        ...

    def forward_train(self，image，target，**kwargs):
        # 返回输出 logits 和损失字典
        ...
    # 如果有更多用于蒸馏的神经网络模块，则重写 get_learnable_parameters 函数。
    # 如果需要获取额外的代价，则重写 get_extra_parameters 函数。
  ...

2. 在 mdistiller/distillers/__init__.py 中的 distiller_dict 中注册该蒸馏器。

3. 在 mdistiller/engines/cfg.py 中注册相应的超参数。

4. 创建一个新的配置文件并进行测试。

引用

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

@article{zhao2022dkd,
  title={解耦知识蒸馏},
  author={Zhao, Borui and Cui, Quan and Song, Renjie and Qiu, Yiyu and Liang, Jiajun},
  journal={arXiv预印本 arXiv:2203.08679},
  year={2022}
}
@article{zhao2023dot,
  title={DOT：面向知识蒸馏的训练器},
  author={Zhao, Borui and Cui, Quan and Song, Renjie and Liang, Jiajun},
  journal={arXiv预印本 arXiv:2307.08436},
  year={2023}
}

许可证

MDistiller 采用 MIT 许可证发布。详情请参阅 LICENSE 文件。

致谢

• 感谢 CRD 和 ReviewKD。我们基于 CRD 的代码库 和 ReviewKD 的代码库 构建了这个库。

• 感谢邱一宇和史毅在旷视科技实习期间的代码贡献。

• 感谢金鑫关于 DKD 的讨论。

tar xvf cifar_teachers.tar -C ./download_ckpts

# 使用默认配置文件训练
python3 tools/train.py --cfg configs/cifar100/dkd/res32x4_res8x4.yaml

# 或通过命令行覆盖参数 (例如修改 Batch Size 和学习率)
python3 tools/train.py --cfg configs/cifar100/dkd/res32x4_res8x4.yaml SOLVER.BATCH_SIZE 128 SOLVER.LR 0.1

# 以 ResNet34 蒸馏到 ResNet18 (DKD 方法) 为例
python3 tools/train.py --cfg configs/imagenet/r34_r18/dkd.yaml