KD-Lib

一个包含知识蒸馏、剪枝和量化等易用方法的 PyTorch 模型压缩库

安装

从源码安装（推荐）

git clone https://github.com/SforAiDl/KD_Lib.git
cd KD_Lib
python setup.py install

从 PyPI 安装

pip install KD-Lib

示例用法

要实现来自 Distilling the Knowledge in a Neural Network 的最基础的知识蒸馏版本，并绘制损失曲线：

import torch
import torch.optim as optim
from torchvision import datasets, transforms
from KD_Lib.KD import VanillaKD

# 这一部分是你定义数据集、数据加载器、模型和优化器的地方

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=False,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

teacher_model = <你的模型>
student_model = <你的模型>

teacher_optimizer = optim.SGD(teacher_model.parameters(), 0.01)
student_optimizer = optim.SGD(student_model.parameters(), 0.01)

# 现在，KD_Lib 就派上用场了

distiller = VanillaKD(teacher_model, student_model, train_loader, test_loader, 
                      teacher_optimizer, student_optimizer)  
distiller.train_teacher(epochs=5, plot_losses=True, save_model=True)    # 训练教师网络
distiller.train_student(epochs=5, plot_losses=True, save_model=True)    # 训练学生网络
distiller.evaluate(teacher=False)                                       # 评估学生网络
distiller.get_parameters()                                              # 一个用于获取教师和学生网络参数数量的实用函数

要使用 Deep Mutual Learning 中的框架，以在线方式训练一组 3 个模型，并将训练细节记录到 TensorBoard 中：

import torch
import torch.optim as optim
from torchvision import datasets, transforms
from KD_Lib.KD import DML
from KD_Lib.models import ResNet18, ResNet50          # 使用 KD_Lib 中打包的模型

# 定义你的数据集、数据加载器、模型和优化器

train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=False,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

student_params = [4, 4, 4, 4, 4]
student_model_1 = ResNet50(student_params, 1, 10)
student_model_2 = ResNet18(student_params, 1, 10)

student_cohort = [student_model_1, student_model_2]

student_optimizer_1 = optim.SGD(student_model_1.parameters(), 0.01)
student_optimizer_2 = optim.SGD(student_model_2.parameters(), 0.01)

student_optimizers = [student_optimizer_1, student_optimizer_2]

# 现在，KD_Lib 就派上用场了

distiller = DML(student_cohort, train_loader, test_loader, student_optimizers, log=True, logdir="./logs")

distiller.train_students(epochs=5)
distiller.evaluate()
distiller.get_parameters()

已实现的方法

部分基准测试结果可在 日志 文件中找到。

论文 / 方法	链接	仓库 (KD_Lib/)
从神经网络中提炼知识	https://arxiv.org/abs/1503.02531	KD/vision/vanilla
通过教师助手改进知识蒸馏	https://arxiv.org/abs/1902.03393	KD/vision/TAKD
关系知识蒸馏	https://arxiv.org/abs/1904.05068	KD/vision/RKD
从噪声教师处蒸馏知识	https://arxiv.org/abs/1610.09650	KD/vision/noisy
更加关注注意力机制	https://arxiv.org/abs/1612.03928	KD/vision/attention
重访知识蒸馏：无教师框架	https://arxiv.org/abs/1909.11723	KD/vision/teacher_free
平均教师是更好的榜样	https://arxiv.org/abs/1703.01780	KD/vision/mean_teacher
通过路径约束优化进行知识蒸馏	https://arxiv.org/abs/1904.09149	KD/vision/RCO
再生神经网络	https://arxiv.org/abs/1805.04770	KD/vision/BANN
准备课程：通过更好的监督改进知识蒸馏	https://arxiv.org/abs/1911.07471	KD/vision/KA
通过知识蒸馏中的噪声协作提升泛化鲁棒性	https://arxiv.org/abs/1910.05057	KD/vision/noisy
将BERT中的任务特定知识蒸馏到简单神经网络	https://arxiv.org/abs/1903.12136	KD/text/BERT2LSTM
深度互学习	https://arxiv.org/abs/1706.00384	KD/vision/DML
彩票假说：寻找稀疏且可训练的神经网络	https://arxiv.org/abs/1803.03635	Pruning/lottery_tickets
通过自我知识蒸馏正则化类别级预测	https://arxiv.org/abs/2003.13964	KD/vision/CSDK

如果您在任何方面觉得 KD-Lib 有用，请引用我们的预印本 :)

@misc{shah2020kdlib,
  title={KD-Lib: 用于知识蒸馏、剪枝和量化的一套 PyTorch 库}, 
  author={Het Shah 和 Avishree Khare 和 Neelay Shah 和 Khizir Siddiqui},
  year={2020},
  eprint={2011.14691},
  archivePrefix={arXiv},
  primaryClass={cs.LG}
}

KD-Lib 快速上手指南

环境准备

系统要求

• Python 3.6 或更高版本
• PyTorch 1.0 或更高版本（建议使用 1.8+）

前置依赖

• torchvision：用于数据集加载和预处理
• numpy：用于数值计算
• matplotlib（可选）：用于绘制损失曲线

推荐使用国内镜像源加速安装，例如：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple

安装步骤

从源码安装（推荐）

git clone https://github.com/SforAiDl/KD_Lib.git
cd KD_Lib
python setup.py install

从 PyPI 安装

pip install KD-Lib

基本使用

知识蒸馏示例（VanillaKD）

import torch
import torch.optim as optim
from torchvision import datasets, transforms
from KD_Lib.KD import VanillaKD

# 数据加载器
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=False,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

# 定义教师模型和学生模型（需自行替换为实际模型）
teacher_model = <your model>
student_model = <your model>

# 定义优化器
teacher_optimizer = optim.SGD(teacher_model.parameters(), 0.01)
student_optimizer = optim.SGD(student_model.parameters(), 0.01)

# 初始化蒸馏器
distiller = VanillaKD(teacher_model, student_model, train_loader, test_loader, 
                      teacher_optimizer, student_optimizer)

# 训练教师模型
distiller.train_teacher(epochs=5, plot_losses=True, save_model=True)

# 训练学生模型
distiller.train_student(epochs=5, plot_losses=True, save_model=True)

# 评估学生模型
distiller.evaluate(teacher=False)

# 查看模型参数量
distiller.get_parameters()

深度互学习示例（DML）

import torch
import torch.optim as optim
from torchvision import datasets, transforms
from KD_Lib.KD import DML
from KD_Lib.models import ResNet18, ResNet50

# 数据加载器
train_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=True,
        download=True,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

test_loader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "mnist_data",
        train=False,
        transform=transforms.Compose(
            [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]
        ),
    ),
    batch_size=32,
    shuffle=True,
)

# 定义学生模型
student_params = [4, 4, 4, 4, 4]
student_model_1 = ResNet50(student_params, 1, 10)
student_model_2 = ResNet18(student_params, 1, 10)

student_cohort = [student_model_1, student_model_2]

# 定义优化器
student_optimizer_1 = optim.SGD(student_model_1.parameters(), 0.01)
student_optimizer_2 = optim.SGD(student_model_2.parameters(), 0.01)

student_optimizers = [student_optimizer_1, student_optimizer_2]

# 初始化蒸馏器
distiller = DML(student_cohort, train_loader, test_loader, student_optimizers, log=True, logdir="./logs")

# 训练学生模型
distiller.train_students(epochs=5)

# 评估
distiller.evaluate()

# 查看模型参数量
distiller.get_parameters()