PyTorch知识蒸馏实战：从零开始训练你的AI模型！

创作时间:

作者:

@小白创作中心

PyTorch知识蒸馏实战：从零开始训练你的AI模型！

引用

github

等

来源

https://github.com/haitongli/knowledge-distillation-pytorch

https://github.com/tyui592/knowledge_distillation

https://blog.csdn.net/shi2xian2wei2/article/details/84570620

https://josehoras.github.io/knowledge-distillation/

https://pytorch.org/examples/

知识蒸馏（Knowledge Distillation）是深度学习领域中一种重要的模型压缩技术，它通过将大型复杂模型（教师模型）的知识传递给小型简单模型（学生模型），实现在保持较高精度的同时减少计算资源消耗。本文将详细介绍如何使用PyTorch框架实现知识蒸馏，并通过MNIST数据集上的实验展示其效果。

知识蒸馏原理

知识蒸馏的核心思想是利用教师模型的"暗知识"（dark knowledge）来指导学生模型的学习。暗知识指的是教师模型在预测时给出的软概率分布，这些分布包含了比简单分类标签更丰富的信息。例如，当一个教师模型在识别手写数字时，即使它确信某个输入是数字"3"，它仍然会给出其他数字（如"8"）的非零概率，这些概率信息就是暗知识。

在传统的监督学习中，模型通常只学习硬标签（one-hot编码），而在知识蒸馏中，学生模型不仅学习硬标签，还学习教师模型的软概率分布。这种学习方式可以看作是一种软标签学习，它能够帮助学生模型更好地理解数据的内在结构，从而在较小的模型规模下达到较高的性能。

PyTorch实现步骤

1. 环境搭建

首先需要确保已经安装了PyTorch及相关依赖。可以使用以下命令安装：

pip install torch torchvision

2. 定义教师模型和学生模型

这里以MNIST数据集为例，教师模型使用一个较大的卷积神经网络，学生模型使用一个较小的全连接网络。

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms

# 教师模型
class TeacherNet(nn.Module):
    def __init__(self):
        super(TeacherNet, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3)
        self.conv2 = nn.Conv2d(32, 64, kernel_size=3)
        self.fc1 = nn.Linear(64*12*12, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = torch.relu(torch.max_pool2d(self.conv1(x), 2))
        x = torch.relu(torch.max_pool2d(self.conv2(x), 2))
        x = x.view(-1, 64*12*12)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

# 学生模型
class StudentNet(nn.Module):
    def __init__(self):
        super(StudentNet, self).__init__()
        self.fc1 = nn.Linear(28*28, 128)
        self.fc2 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 28*28)
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

3. 蒸馏损失函数

知识蒸馏的关键在于设计合适的损失函数。通常使用交叉熵损失结合蒸馏损失，其中蒸馏损失通过软概率分布计算。

def distillation_loss(y, labels, teacher_scores, T, alpha):
    return nn.KLDivLoss()(F.log_softmax(y/T, dim=1),
                          F.softmax(teacher_scores/T, dim=1)) * (T*T * 2.0 * alpha) + \
           F.cross_entropy(y, labels) * (1. - alpha)

4. 训练过程

完整的训练代码如下：

def train_student(student, teacher, device, train_loader, optimizer, epoch, T, alpha):
    student.train()
    teacher.eval()
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        output = student(data)
        with torch.no_grad():
            teacher_output = teacher(data)
        loss = distillation_loss(output, target, teacher_output, T, alpha)
        loss.backward()
        optimizer.step()
        if batch_idx % 100 == 0:
            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                epoch, batch_idx * len(data), len(train_loader.dataset),
                100. * batch_idx / len(train_loader), loss.item()))

def test(model, device, test_loader):
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            test_loss += F.cross_entropy(output, target, reduction='sum').item()
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    print('\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
        test_loss, correct, len(test_loader.dataset),
        100. * correct / len(test_loader.dataset)))

# 主程序
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
teacher = TeacherNet().to(device)
student = StudentNet().to(device)
optimizer = optim.Adam(student.parameters(), lr=0.001)

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

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

# 预训练教师模型
teacher.load_state_dict(torch.load('teacher_model.pth'))

# 开始训练学生模型
for epoch in range(1, 11):
    train_student(student, teacher, device, train_loader, optimizer, epoch, T=20, alpha=0.7)
    test(student, device, test_loader)