4. 中间层特征蒸馏
任务目标
本案例实现中间层特征蒸馏：训练一个教师网络（较大模型）后，提取其某一中间层的特征图，设计蒸馏损失使学生网络（较小模型）的对应中间层特征图向教师网络对齐，从而提升学生网络在图像分类任务上的性能。
环境与依赖
需要 Python 3.8+、PyTorch 1.10+、torchvision。数据集使用 CIFAR-10，代码会自动下载。
核心实现
完整代码如下，包含教师网络、学生网络、特征蒸馏损失函数以及训练/验证主流程。教师网络与学生网络均使用简单CNN，教师网络更宽更深。蒸馏发生在第2个卷积块之后，通过一个1×1卷积将学生特征图通道数对齐到教师特征图通道数，然后计算MSE损失。
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
# ---------- 网络定义 ----------
class ConvBlock(nn.Module):
def __init__(self, in_c, out_c, pool=True):
super().__init__()
self.conv1 = nn.Conv2d(in_c, out_c, 3, padding=1)
self.bn1 = nn.BatchNorm2d(out_c)
self.conv2 = nn.Conv2d(out_c, out_c, 3, padding=1)
self.bn2 = nn.BatchNorm2d(out_c)
self.pool = nn.MaxPool2d(2) if pool else nn.Identity()
def forward(self, x):
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
return self.pool(x)
class TeacherNet(nn.Module):
def __init__(self):
super().__init__()
self.block1 = ConvBlock(3, 64)
self.block2 = ConvBlock(64, 128)
self.block3 = ConvBlock(128, 256)
self.fc = nn.Linear(256 * 4 * 4, 10)
def forward(self, x):
x = self.block1(x)
x = self.block2(x)
x = self.block3(x)
x = x.view(x.size(0), -1)
return self.fc(x)
class StudentNet(nn.Module):
def __init__(self):
super().__init__()
self.block1 = ConvBlock(3, 32)
self.block2 = ConvBlock(32, 64)
self.block3 = ConvBlock(64, 128)
self.fc = nn.Linear(128 * 4 * 4, 10)
def forward(self, x):
x = self.block1(x)
x = self.block2(x)
x = self.block3(x)
x = x.view(x.size(0), -1)
return self.fc(x)
# ---------- 特征蒸馏损失 ----------
class FeatureDistillLoss(nn.Module):
def __init__(self, student_channel=64, teacher_channel=128):
super().__init__()
# 将学生特征图通道对齐到教师特征图通道
self.align = nn.Conv2d(student_channel, teacher_channel, 1)
def forward(self, s_feat, t_feat):
s_feat = self.align(s_feat)
return F.mse_loss(s_feat, t_feat.detach())
# ---------- 数据加载 ----------
transform = transforms.Compose([
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2470, 0.2435, 0.2616))
])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
train_loader = DataLoader(trainset, batch_size=64, shuffle=True, num_workers=2)
test_loader = DataLoader(testset, batch_size=64, shuffle=False, num_workers=2)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# ---------- 训练教师网络 ----------
def train_teacher():
teacher = TeacherNet().to(device)
opt = optim.Adam(teacher.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
for epoch in range(10):
teacher.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
opt.zero_grad()
loss = criterion(teacher(x), y)
loss.backward()
opt.step()
torch.save(teacher.state_dict(), 'teacher.pth')
return teacher
# ---------- 特征提取钩子 ----------
def get_feature_hook(model, layer_name):
features = {}
def hook_fn(module, input, output):
features[layer_name] = output
for name, module in model.named_modules():
if name == layer_name:
module.register_forward_hook(hook_fn)
return features
# ---------- 训练学生网络（含特征蒸馏） ----------
def train_student_with_distill(teacher):
teacher.eval()
student = StudentNet().to(device)
opt = optim.Adam(student.parameters(), lr=1e-3)
ce_loss = nn.CrossEntropyLoss()
distill_loss_fn = FeatureDistillLoss().to(device)
# 注册钩子：block2输出（即第2个卷积块后的特征图）
teacher_feats = get_feature_hook(teacher, 'block2')
student_feats = get_feature_hook(student, 'block2')
alpha = 0.5 # 蒸馏损失权重
for epoch in range(10):
student.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
opt.zero_grad()
out = student(x)
loss_ce = ce_loss(out, y)
# 教师前向传播（不计算梯度）
with torch.no_grad():
teacher(x)
loss_distill = distill_loss_fn(student_feats['block2'], teacher_feats['block2'])
loss = loss_ce + alpha * loss_distill
loss.backward()
opt.step()
return student
# ---------- 验证 ----------
def evaluate(model):
model.eval()
correct = 0
total = 0
with torch.no_grad():
for x, y in test_loader:
x, y = x.to(device), y.to(device)
pred = model(x).argmax(dim=1)
correct += (pred == y).sum().item()
total += y.size(0)
return correct / total
if __name__ == '__main__':
# 1. 训练教师网络
teacher = train_teacher()
teacher_acc = evaluate(teacher)
print(f'Teacher accuracy: {teacher_acc:.4f}')
# 2. 训练学生网络（带特征蒸馏）
student = train_student_with_distill(teacher)
student_acc = evaluate(student)
print(f'Student (with distill) accuracy: {student_acc:.4f}')
# 3. 对比：训练不带蒸馏的学生网络
student_no_distill = StudentNet().to(device)
opt = optim.Adam(student_no_distill.parameters(), lr=1e-3)
criterion = nn.CrossEntropyLoss()
for epoch in range(10):
student_no_distill.train()
for x, y in train_loader:
x, y = x.to(device), y.to(device)
opt.zero_grad()
loss = criterion(student_no_distill(x), y)
loss.backward()
opt.step()
student_no_distill_acc = evaluate(student_no_distill)
print(f'Student (no distill) accuracy: {student_no_distill_acc:.4f}')运行方式
将上述代码保存为 feature_distill.py，直接运行：
python feature_distill.py首次运行会自动下载CIFAR-10数据集到 ./data 目录。代码依次执行：训练教师网络（10轮）→ 训练带蒸馏的学生网络（10轮）→ 训练不带蒸馏的学生网络（10轮），最后打印三个模型的测试集准确率。
验证方法与判定标准
运行结束后观察终端输出的三个准确率。判定通过的依据：
- 教师网络准确率明显高于不带蒸馏的学生网络（通常差距在5%以上）；
- 带蒸馏的学生网络准确率高于不带蒸馏的学生网络，说明中间层特征蒸馏确实带来了性能提升。
由于随机种子未固定，每次运行准确率会有小幅波动，但上述相对关系在大多数情况下成立。若带蒸馏的学生准确率未超过不带蒸馏版本，可尝试增大 alpha 至0.7或增加训练轮数至20轮。
关键说明
本案例的蒸馏点选择在 block2 输出（即第2个卷积块后的特征图），教师特征图尺寸为 128×8×8，学生为 64×8×8。通过 FeatureDistillLoss 中的1×1卷积将学生通道从64对齐到128，再计算MSE损失。教师特征图使用 detach() 阻断梯度回传，确保只优化学生网络。若需蒸馏其他层，只需修改 get_feature_hook 中的层名并调整 FeatureDistillLoss 的通道参数即可。