生成实现知识蒸馏的代码

以下是一个简单的示例代码，用于实现知识蒸馏：

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

# 定义教师模型和学生模型
teacher_model = TeacherModel()
student_model = StudentModel()

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(student_model.parameters(), lr=0.001)

# 加载数据集
train_data = MyDataset(...)
train_loader = DataLoader(train_data, batch_size=32, shuffle=True)

# 训练学生模型
for epoch in range(10):
    student_model.train()
    running_loss = 0.0
    for data in train_loader:
        inputs, labels = data
        optimizer.zero_grad()
        
        # 前向传播
        with torch.no_grad():
            teacher_outputs = teacher_model(inputs)
        student_outputs = student_model(inputs)
        
        # 计算损失
        loss = criterion(student_outputs, teacher_outputs)
        
        # 反向传播和优化
        loss.backward()
        optimizer.step()
        
        # 统计损失
        running_loss += loss.item() * inputs.size(0)
    
    # 输出每个 epoch 的统计结果
    epoch_loss = running_loss / len(train_data)
    print('Epoch {}/{} loss: {:.4f}'.format(epoch+1, 10, epoch_loss))

在上面的代码中，我们首先定义了教师模型和学生模型，并且使用均方误差损失函数作为知识蒸馏的目标。然后使用 Adam 优化器来训练学生模型。在前向传播时，我们使用教师模型的输出作为目标输出，并计算学生模型的输出与目标输出之间的损失。最后反向传播并更新学生模型的参数。在每个 epoch 结束时，我们输出该 epoch 的平均损失。