ResNet18模型蒸馏实战：云端双GPU对比实验

引言

作为一名研究生，当你需要在论文中验证模型压缩算法的效果时，可能会遇到这样的困境：本地只有单张GPU显卡，而实验需要对比不同配置下的模型性能。特别是像ResNet18这样的经典网络，在模型蒸馏（知识蒸馏）过程中，教师模型和学生模型的并行训练往往需要多GPU环境支持。

本文将带你用最简单的方式，在云端快速搭建双GPU实验环境，完成ResNet18模型蒸馏的完整对比实验。整个过程就像在实验室借用了一台高性能工作站，但成本和时间消耗却低得多。我们会从环境准备开始，一步步完成数据准备、模型训练、蒸馏实现和结果对比，所有代码都可以直接复制运行。

1. 环境准备与镜像选择

首先我们需要一个预装好PyTorch和CUDA的深度学习环境。这里推荐使用CSDN星图镜像广场中的PyTorch官方镜像，它已经预装了：

PyTorch 1.12+ 和 torchvision
CUDA 11.6 和 cuDNN
常用的Python数据科学包（numpy, pandas等）

选择这个镜像的原因很简单：它就像是一个已经配好所有调料的厨房，开箱即用。特别是对于多GPU支持，官方镜像已经做好了底层配置，我们不需要自己折腾NCCL通信这些复杂的东西。

启动实例时，记得选择至少2块GPU的计算规格。在CSDN算力平台上，你可以直接搜索"PyTorch"找到这个镜像，点击部署后选择"2*GPU"的配置即可。

2. 数据准备与模型加载

环境准备好后，我们先准备实验用的数据集和模型。这里以CIFAR-10为例，因为它大小适中，适合快速实验：

import torch import torchvision import torchvision.transforms as transforms # 数据预处理 transform = transforms.Compose([ transforms.Resize(224), # ResNet18的标准输入尺寸 transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) ]) # 加载CIFAR-10数据集 trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform) trainloader = torch.utils.data.DataLoader(trainset, batch_size=256, shuffle=True, num_workers=2) # 加载预训练的ResNet18模型 teacher_model = torchvision.models.resnet18(pretrained=True) student_model = torchvision.models.resnet18(pretrained=False) # 学生模型不加载预训练权重

这里有个小技巧：虽然CIFAR-10的原始尺寸是32x32，但我们调整到224x224以适应ResNet18的标准输入。这不会影响实验对比的公平性，因为教师模型和学生模型都会使用相同的输入尺寸。

3. 多GPU并行配置

现在来到关键步骤——配置多GPU训练。PyTorch让这个过程变得非常简单：

# 检查可用GPU数量 device_count = torch.cuda.device_count() print(f"可用GPU数量: {device_count}") # 将模型分布到多个GPU上 if device_count > 1: teacher_model = torch.nn.DataParallel(teacher_model) student_model = torch.nn.DataParallel(student_model) # 移动到GPU device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") teacher_model.to(device) student_model.to(device)

DataParallel是PyTorch提供的包装器，它会自动将输入数据分割并分发到各个GPU，然后收集梯度进行同步更新。这就像有一个智能的调度员，帮你把工作分配给多个工人，最后汇总结果。

4. 知识蒸馏实现

知识蒸馏的核心思想是让学生模型不仅学习真实标签，还要模仿教师模型的"软标签"（soft targets）。下面是蒸馏损失函数的实现：

import torch.nn as nn import torch.nn.functional as F class DistillationLoss(nn.Module): def __init__(self, temperature=4.0, alpha=0.7): super().__init__() self.temperature = temperature self.alpha = alpha self.kl_div = nn.KLDivLoss(reduction='batchmean') def forward(self, student_logits, teacher_logits, labels): # 计算硬损失（常规交叉熵） hard_loss = F.cross_entropy(student_logits, labels) # 计算软损失（KL散度） soft_loss = self.kl_div( F.log_softmax(student_logits/self.temperature, dim=1), F.softmax(teacher_logits/self.temperature, dim=1) ) * (self.temperature ** 2) # 组合损失 return self.alpha * soft_loss + (1 - self.alpha) * hard_loss

这里的temperature参数控制着教师模型输出的"软化"程度。温度越高，概率分布越平滑，学生模型能学到更多类别间的关系。alpha参数则平衡了硬标签和软标签的重要性。

5. 训练过程与对比实验

现在我们可以开始训练了。为了对比单GPU和双GPU的效果，我们分别运行两个实验：

def train_model(model, criterion, optimizer, epochs=10, desc=""): model.train() for epoch in range(epochs): running_loss = 0.0 for i, (inputs, labels) in enumerate(trainloader, 0): inputs, labels = inputs.to(device), labels.to(device) optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, labels) loss.backward() optimizer.step() running_loss += loss.item() if i % 50 == 49: print(f"[{desc}] Epoch {epoch+1}, Batch {i+1}: loss {running_loss/50:.3f}") running_loss = 0.0 # 教师模型训练（单GPU） teacher_optimizer = torch.optim.SGD(teacher_model.parameters(), lr=0.01, momentum=0.9) train_model(teacher_model, nn.CrossEntropyLoss(), teacher_optimizer, epochs=5, desc="Teacher") # 学生模型蒸馏训练（双GPU） distill_criterion = DistillationLoss(temperature=4.0, alpha=0.7) student_optimizer = torch.optim.SGD(student_model.parameters(), lr=0.01, momentum=0.9) train_model(student_model, distill_criterion, student_optimizer, epochs=10, desc="Student")

在实际论文实验中，你还需要：

记录每个epoch的训练时间，比较单GPU和双GPU的速度差异
在验证集上测试模型准确率，比较教师模型和学生模型的性能差距
尝试不同的温度参数和alpha值，找到最优的蒸馏配置

6. 实验结果分析与可视化

实验完成后，我们可以用Matplotlib绘制一些对比图表：

import matplotlib.pyplot as plt # 假设我们已经记录了以下实验数据 single_gpu_time = [120, 118, 119, 117, 120] # 单GPU每个epoch的时间(秒) dual_gpu_time = [65, 64, 66, 65, 64] # 双GPU每个epoch的时间(秒) teacher_acc = [0.75, 0.82, 0.85, 0.86, 0.87] # 教师模型验证准确率 student_acc = [0.72, 0.80, 0.83, 0.85, 0.86] # 学生模型验证准确率 # 绘制训练时间对比 plt.figure(figsize=(10, 4)) plt.subplot(1, 2, 1) plt.plot(single_gpu_time, label='Single GPU') plt.plot(dual_gpu_time, label='Dual GPU') plt.title('Training Time per Epoch') plt.ylabel('Seconds') plt.legend() # 绘制准确率对比 plt.subplot(1, 2, 2) plt.plot(teacher_acc, label='Teacher Model') plt.plot(student_acc, label='Student Model') plt.title('Validation Accuracy') plt.ylabel('Accuracy') plt.legend() plt.tight_layout() plt.show()

从实验结果中，你通常会发现：