知识蒸馏 pytorch代码

知识蒸馏是一种使用已训练好的大模型（教师模型）来指导训练小模型（学生模型）的技术，从而提高小模型的性能和泛化能力。下面是一个使用 PyTorch 实现知识蒸馏的代码示例：

import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision.datasets import CIFAR10
from torchvision.transforms import ToTensor
from torchvision.models import resnet18

# 定义教师模型
teacher_model = resnet18(pretrained=True)
teacher_model.eval()

# 定义学生模型
student_model = resnet18()
student_model.train()

# 定义损失函数
criterion = nn.KLDivLoss()

# 定义优化器
optimizer = optim.SGD(student_model.parameters(), lr=0.001, momentum=0.9)

# 加载 CIFAR-10 数据集
train_dataset = CIFAR10(root='./data', train=True, transform=ToTensor())
train_loader = DataLoader(train_dataset, batch_size=128, shuffle=True)

# 训练学生模型
for epoch in range(10):
    running_loss = 0.0
    for inputs, targets in train_loader:
        # 使用教师模型生成 soft targets
        with torch.no_grad():
            soft_targets = nn.functional.softmax(teacher_model(inputs), dim=1)

        # 使用学生模型生成 outputs
        outputs = student_model(inputs)

        # 计算损失函数
        loss = criterion(nn.functional.log_softmax(outputs, dim=1), soft_targets)

        # 反向传播和优化
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 输出当前损失函数值
        running_loss += loss.item()
    print(f'Epoch {epoch+1}, loss: {running_loss/len(train_loader)}')

在上面的代码中，我们首先定义了一个已经训练好的 ResNet-18 模型作为教师模型，并将其设为 eval 模式。然后我们定义了一个未训练的 ResNet-18 模型作为学生模型，并将其设为 train 模式。接着，我们定义了一个 Kullback-Leibler 散度损失函数作为我们的损失函数，并定义了一个随机梯度下降优化器。最后，我们加载 CIFAR-10 数据集，并训练学生模型。

在训练过程中，我们使用教师模型来生成 soft targets（也就是概率分布），并将其作为标签来训练学生模型。这样做的目的是让学生模型学习教师模型的知识。具体地，我们首先使用教师模型对输入数据进行前向传播，并计算出其在各个类别上的概率分布。然后，我们使用学生模型对输入数据进行前向传播，并计算出其在各个类别上的概率分布。最后，我们使用 Kullback-Leibler 散度损失函数来计算学生模型的输出概率分布和教师模型的输出概率分布之间的差异，并利用反向传播算法和随机梯度下降优化器来更新学生模型的参数。