知识蒸馏代码实现(以MNIST手写数字体为例,自定义MLP网络做为教师和学生网络)

dataloader_tools.py

import torchvision
from torchvision import transforms
from torch.utils.data import DataLoaderdef load_data():# 载入MNIST训练集train_dataset = torchvision.datasets.MNIST(root = "../datasets/",train=True,transform=transforms.ToTensor(),download=True)# 载入MNIST测试集test_dataset = torchvision.datasets.MNIST(root = "../datasets/",train=False,transform=transforms.ToTensor(),download=True)# 生成训练集和测试集的dataloadertrain_dataloader = DataLoader(dataset=train_dataset,batch_size=12,shuffle=True)test_dataloader = DataLoader(dataset=test_dataset,batch_size=12,shuffle=False)return train_dataloader,test_dataloader

models.py

import torch
from torch import nn
# 教师模型
class TeacherModel(nn.Module):def __init__(self,in_channels=1,num_classes=10):super(TeacherModel,self).__init__()self.relu = nn.ReLU()self.fc1 = nn.Linear(784,1200)self.fc2 = nn.Linear(1200,1200)self.fc3 = nn.Linear(1200,num_classes)self.dropout = nn.Dropout(p=0.5) #p=0.5是丢弃该层一半的神经元.def forward(self,x):x = x.view(-1,784)x = self.fc1(x)x = self.dropout(x)x = self.relu(x)x = self.fc2(x)x = self.dropout(x)x = self.relu(x)x = self.fc3(x)return xclass StudentModel(nn.Module):def __init__(self,in_channels=1,num_classes=10):super(StudentModel,self).__init__()self.relu = nn.ReLU()self.fc1 = nn.Linear(784,20)self.fc2 = nn.Linear(20,20)self.fc3 = nn.Linear(20,num_classes)def forward(self,x):x = x.view(-1,784)x = self.fc1(x)x = self.relu(x)x = self.fc2(x)x = self.relu(x)x = self.fc3(x)return x

train_tools.py

from torch import nn
import time
import torch
import tqdm
import torch.nn.functional as Fdef train(epochs, model, model_name, lr,train_dataloader,test_dataloader,device):# ----------------------开始计时-----------------------------------start_time = time.time()# 设置参数开始训练best_acc, best_epoch = 0, 0criterion = nn.CrossEntropyLoss()optimizer = torch.optim.Adam(model.parameters(), lr=lr)for epoch in range(epochs):model.train()# 训练集上训练模型权重for data, targets in tqdm.tqdm(train_dataloader):# 把数据加载到GPU上data = data.to(device)targets = targets.to(device)# 前向传播preds = model(data)loss = criterion(preds, targets)# 反向传播optimizer.zero_grad()loss.backward()optimizer.step()# 测试集上评估模型性能model.eval()num_correct = 0num_samples = 0with torch.no_grad():for x, y in test_dataloader:x = x.to(device)y = y.to(device)preds = model(x)predictions = preds.max(1).indices  # 返回每一行的最大值和该最大值在该行的列索引num_correct += (predictions == y).sum()num_samples += predictions.size(0)acc = (num_correct / num_samples).item()if acc > best_acc:best_acc = accbest_epoch = epoch# 保存模型最优准确率的参数torch.save(model.state_dict(), f"../weights/{model_name}_best_acc_params.pth")model.train()print('Epoch:{}\t Accuracy:{:.4f}'.format(epoch + 1, acc),f'loss={loss}')print(f'最优准确率的epoch为{best_epoch},值为:{best_acc},最优参数已经保存到:weights/{model_name}_best_acc_params.pth')# -------------------------结束计时------------------------------------end_time = time.time()run_time = end_time - start_time# 将输出的秒数保留两位小数if int(run_time) < 60:print(f'训练用时为:{round(run_time, 2)}s')else:print(f'训练用时为:{round(run_time / 60, 2)}minutes')def distill_train(epochs,teacher_model,student_model,model_name,train_dataloader,test_dataloader,alpha,lr,temp,device):# -------------------------------------开始计时--------------------------------start_time = time.time()# 定以损失函数hard_loss = nn.CrossEntropyLoss()soft_loss = nn.KLDivLoss(reduction="batchmean")# 定义优化器optimizer = torch.optim.Adam(student_model.parameters(), lr=lr)best_acc,best_epoch = 0,0for epoch in range(epochs):student_model.train()# 训练集上训练模型权重for data,targets in tqdm.tqdm(train_dataloader):# 把数据加载到GPU上data = data.to(device)targets = targets.to(device)# 教师模型预测with torch.no_grad():teacher_preds = teacher_model(data)# 学生模型预测student_preds = student_model(data)# 计算hard_lossstudent_hard_loss = hard_loss(student_preds,targets)# 计算蒸馏后的预测结果及soft_lossditillation_loss = soft_loss(F.softmax(student_preds/temp,dim=1),F.softmax(teacher_preds/temp,dim=1))# 将hard_loss和soft_loss加权求和loss = temp * temp * alpha * student_hard_loss + (1-alpha)*ditillation_loss# 反向传播,优化权重optimizer.zero_grad()loss.backward()optimizer.step()#测试集上评估模型性能student_model.eval()num_correct = 0num_samples = 0with torch.no_grad():for x,y in test_dataloader:x = x.to(device)y = y.to(device)preds = student_model(x)predictions = preds.max(1).indices #返回每一行的最大值和该最大值在该行的列索引num_correct += (predictions ==y).sum()num_samples += predictions.size(0)acc = (num_correct/num_samples).item()if acc>best_acc:best_acc = accbest_epoch = epoch# 保存模型最优准确率的参数torch.save(student_model.state_dict(),f"../weights/{model_name}_best_acc_params.pth")student_model.train()print('Epoch:{}\t Accuracy:{:.4f}'.format(epoch+1,acc))print(f'student_hard_loss={student_hard_loss},ditillation_loss={ditillation_loss},loss={loss}')print(f'最优准确率的epoch为{best_epoch},值为:{best_acc},')# --------------------------------结束计时----------------------------------end_time = time.time()run_time = end_time - start_time# 将输出的秒数保留两位小数if int(run_time) < 60:print(f'训练用时为:{round(run_time, 2)}s')else:print(f'训练用时为:{round(run_time / 60, 2)}minutes')

训练教师网络

import torch
from torchinfo import summary #用来可视化的
import models
import dataloader_tools
import train_tools# 设置随机数种子
torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 使用cuDNN加速卷积运算
torch.backends.cudnn.benchmark = True# 载入MNIST训练集和测试集
train_dataloader,test_dataloader = dataloader_tools.load_data()# 定义教师模型
model = models.TeacherModel()
model = model.to(device)
# 打印模型的参数
summary(model)# 定义参数并开始训练
epochs = 10
lr = 1e-4
model_name = 'teacher'
train_tools.train(epochs,model,model_name,lr,train_dataloader,test_dataloader,device)

最优准确率的epoch为9,值为:0.9868999719619751

用非蒸馏的方法训练学生网络

import torch
from torchinfo import summary #用来可视化的
import dataloader_tools
import models
import train_tools# 设置随机数种子
torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 使用cuDNN加速卷积运算
torch.backends.cudnn.benchmark = True# 生成训练集和测试集的dataloader
train_dataloader,test_dataloader = dataloader_tools.load_data()# 从头训练学生模型
model = models.StudentModel()
model = model.to(device)
# 查看模型参数
print(summary(model))# 定义参数并开始训练
epochs = 10
lr = 1e-4
model_name = 'student'
train_tools.train(epochs, model, model_name, lr,train_dataloader,test_dataloader,device)

最优准确率的epoch为9,准确率为:0.9382999539375305,最优参数已经保存到:weights/student_best_acc_params.pth
训练用时为:1.74minutes

用知识蒸馏的方法训练student model

import torch
import train_tools
import models
import dataloader_tools# 设置随机数种子
torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 使用cuDNN加速卷积运算
torch.backends.cudnn.benchmark = True# 加载数据
train_dataloader,test_dataloader = dataloader_tools.load_data()# 加载训练好的teacher model
teacher_model = models.TeacherModel()
teacher_model = teacher_model.to(device)
teacher_model.load_state_dict(torch.load('../weights/teacher_best_acc_params.pth'))
teacher_model.eval()# 准备新的学生模型
student_model = models.StudentModel()
student_model = student_model.to(device)
student_model.train()# 开始训练
lr = 0.0001
epochs = 20
alpha = 0.3 # hard_loss权重
temp = 7 # 蒸馏温度
model_name = 'distill_student_loss'
# 调用train_tools中的
train_tools.distill_train(epochs,teacher_model,student_model,model_name,train_dataloader,test_dataloader,alpha,lr,temp,device)

最优准确率的epoch为9,值为:0.9204999804496765,
训练用时为:2.14minutes

loss改为：

# temp的平方乘在student_hard_loss
loss = temp * temp * alpha * student_hard_loss + (1 - alpha) * ditillation_loss

最优准确率的epoch为9,值为:0.9336999654769897,
训练用时为:2.12minutes

loss改为：

# temp的平方乘ditillation_loss
loss = alpha * student_hard_loss + temp * temp * (1 - alpha) * ditillation_loss

最优准确率的epoch为9,值为:0.9176999926567078,
训练用时为:2.09minutes

上面的几种loss，蒸馏损失都出现了负数的情况。不太对劲。

其它开源的知识蒸馏算法如下：

open-mmlab开源的工具箱包含知识蒸馏算法

mmrazor

github.com/open-mmlab/mmrazor

NAS：神经架构搜索
剪枝：Pruning
KD: 知识蒸馏
Quantization: 量化

自定义知识蒸馏算法：

mmdeploy

可以把算法部署到一些厂商支持的中间格式，如ONNX，tensorRT等。

HobbitLong的RepDistiller

github.com/HobbitLong/RepDistiller

里面有12种最新的知识蒸馏算法。

蒸馏网络可以应用于同一种模型，将大的学习的知识蒸馏到小的上面。
如下将resnet100做教师网络，resnet32做学生网络。

将一种模型迁移到另一种模型上。如vgg13做教师网络，mobilNetv2做学生网络：