PyTorch实战：MNIST知识蒸馏全解析与实现指南

1次阅读

没有评论

共计 2937 个字符，预计需要花费 8 分钟才能阅读完成。

知识蒸馏（Knowledge Distillation）是一种模型压缩技术，其核心思想是通过训练一个轻量级的“学生模型”来模仿一个更复杂的“教师模型”的行为。这种方法的价值主要体现在以下几个方面：

模型压缩：学生模型通常比教师模型更小、更快，适合部署在资源受限的设备上。
性能提升：学生模型通过模仿教师模型的“软标签”（概率分布），往往能比直接训练表现得更好。
泛化能力：教师模型提供的额外信息（如类别间的关系）可以帮助学生模型更好地泛化。

MNIST 是一个手写数字识别数据集，包含 60,000 张训练图像和 10,000 张测试图像。每张图像是 28×28 的灰度图，标签为 0 - 9 的数字。预处理通常包括以下步骤：

归一化：将像素值从 [0, 255] 缩放到[0, 1]。
标准化：进一步调整均值为 0，标准差为 1。
数据增强：如随机旋转或平移，以增加数据多样性。

教师模型通常是一个复杂的网络，而学生模型则是一个轻量级网络。以下是两种模型的典型架构：

教师模型：
输入层：28×28
卷积层：Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
激活函数：ReLU
池化层：MaxPool2d(kernel_size=2, stride=2)
全连接层：Linear(141432, 10)
学生模型：
输入层：28×28
全连接层：Linear(28*28, 128)
激活函数：ReLU
全连接层：Linear(128, 10)

以下是知识蒸馏的完整实现代码，包括损失函数设计和训练流程：

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 定义教师模型
class TeacherModel(nn.Module):
    def __init__(self):
        super(TeacherModel, self).__init__()
        self.conv1 = nn.Conv2d(1, 32, kernel_size=3, stride=1, padding=1)
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc = nn.Linear(14*14*32, 10)

    def forward(self, x):
        x = self.pool(torch.relu(self.conv1(x)))
        x = x.view(-1, 14*14*32)
        x = self.fc(x)
        return x

# 定义学生模型
class StudentModel(nn.Module):
    def __init__(self):
        super(StudentModel, 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

# 定义知识蒸馏损失函数
def distillation_loss(student_output, teacher_output, temperature):
    soft_teacher = torch.softmax(teacher_output / temperature, dim=1)
    soft_student = torch.log_softmax(student_output / temperature, dim=1)
    return nn.KLDivLoss(reduction='batchmean')(soft_student, soft_teacher)

# 数据加载
transform = transforms.Compose([transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transform)
test_dataset = datasets.MNIST(root='./data', train=False, transform=transform)

train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=64, shuffle=False)

# 初始化模型和优化器
teacher = TeacherModel()
student = StudentModel()
optimizer = optim.Adam(student.parameters(), lr=0.001)

# 训练教师模型
# ...（省略教师模型训练代码）# 知识蒸馏训练
for epoch in range(10):
    for batch_idx, (data, target) in enumerate(train_loader):
        optimizer.zero_grad()
        teacher_output = teacher(data)
        student_output = student(data)

        # 计算损失
        hard_loss = nn.CrossEntropyLoss()(student_output, target)
        soft_loss = distillation_loss(student_output, teacher_output, temperature=5)
        total_loss = hard_loss + soft_loss

        total_loss.backward()
        optimizer.step()

经过知识蒸馏训练后，学生模型的性能通常会比直接训练有所提升。以下是可能的性能对比：