知识蒸馏属于模型的压缩一种方法，但其实这种方法又属于一种伪压缩，是将一个性能较好的teacher网络“压缩”进一个性能较差的student网络中，或者是可类似于在teacher的指导下让student进行学习进而提高性能。

知识蒸馏是一种思想，并不像其他压缩方法有现成的库，因此对于实际需求与场景需要自己去实现。蒸馏也分为“离线”蒸馏与“在线”蒸馏。前者是建立T-S进行KD训练，而后者可以说是一种自学习，让student自己做自己的teacher。

同时蒸馏还分为逻辑蒸馏和特征蒸馏，前者是在两个网络最终输出部分建立loss关系，而后者是在网络中间的某些特征部分建立loss进行蒸馏。

本文是以手写数字为例，teacher选用的resnet18，student选用的resnet50【大家可能会想resnet50比resnet18强啊，为啥resnet50是student，这是因为我在实际测试的时候发现在手写数字这个数据上resnet18的准确率比resnet50高，猜测是因为在低分辨率下resnet50虽然loss在下降，但由于网络较深，特征丢失也明显，网络退化较明显】。当然这里你也可以尝试resnet做teacher，mobilnet做student【我这样训练了一下发现对mobilnet提升变化不大】

注：这里不做模型和蒸馏改进，仅仅是给大家展示一下效果，至于更细化的蒸馏有兴趣的可以自己去研究。【有关目标检测方面的KD 训练，我将会在明年以后推出】

目录

teacher train代码

student未KD 训练 

KD train代码

KD_loss代码：

完整代码

teacher train代码

参数说明：

teacher_model:选用的teacher网络

train_loader:训练集

test_loader:测试集

loss_func:损失函数

epochs:训练迭代数

def teacher_train(teacher_model, train_loader, test_loader, loss_func, epochs):teacher_model.train()teacher_model.cuda()# trainfor i in range(epochs):for data, label in train_loader:data = data.to(device)label = label.to(device)output = teacher_model(data)loss = loss_func(output, label)optimizer_teacher.zero_grad()loss.backward()optimizer_teacher.step()print("loss: ", loss)# evalcorrect = 0teacher_model.eval()teacher_model.cuda()for test_data, test_label in test_loader:test_data = test_data.to(device)test_label = test_label.to(device)with torch.no_grad():output = teacher_model(test_data)# acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))# print("teacher acc: ", acc)_, pred = torch.max(output, dim=1)correct += float(torch.sum(pred == test_label))print('test_acc:{}'.format(correct / len(test_dataset)))return teacher_model

训练结果(我只训练了5轮)： 

teacher model train
loss:  tensor(0.0891, device='cuda:0', grad_fn=)
test_acc:0.9845
loss:  tensor(0.0132, device='cuda:0', grad_fn=)
test_acc:0.9865
loss:  tensor(0.0019, device='cuda:0', grad_fn=)
test_acc:0.9909
loss:  tensor(0.0042, device='cuda:0', grad_fn=)
test_acc:0.9909
loss:  tensor(0.0034, device='cuda:0', grad_fn=)
test_acc:0.9917
teacher model trained finished!

student未KD 训练 

参数说明：

student_model:选用的student网络

train_loader:训练集

test_loader:测试集

loss_func:损失函数

epochs:训练迭代数

def student_train(student_model, train_loader, test_loader, loss_func, epochs):student_model.train()student_model.cuda()# trainfor i in range(epochs):for data, label in train_loader:data = data.to(device)label = label.to(device)output = student_model(data)loss = loss_func(output, label)optimizer_student.zero_grad()loss.backward()optimizer_student.step()print("student loss: ", loss)# evalcorrect = 0student_model.eval()student_model.cuda()for test_data, test_label in test_loader:test_data = test_data.to(device)test_label = test_label.to(device)with torch.no_grad():output = student_model(test_data)# acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))# print("teacher acc: ", acc)_, pred = torch.max(output, dim=1)correct += float(torch.sum(pred == test_label))print('student test_acc:{}'.format(correct / len(test_dataset)))

没有KD train的效果如下： 

student model ready train
student loss:  tensor(0.1876, device='cuda:0', grad_fn=)
student test_acc:0.9588
student loss:  tensor(0.0219, device='cuda:0', grad_fn=)
student test_acc:0.9737
student loss:  tensor(0.0588, device='cuda:0', grad_fn=)
student test_acc:0.9812
student loss:  tensor(0.0024, device='cuda:0', grad_fn=)
student test_acc:0.9853
student loss:  tensor(0.0022, device='cuda:0', grad_fn=)
student test_acc:0.9814student model trained finished!

KD train代码

参数说明：

teacher_model:为已经训练好的teacher

student_model:待KD的student网络

train_loader:训练集

test_loader:测试集

def KD_train(teacher_model, student_model, train_loader, test_loader,loss_func, epochs):teacher_model.eval()student_model.train()student_model.cuda()HL = nn.CrossEntropyLoss()for i in range(epochs):for data, labels in train_loader:data = data.to(device)labels = labels.to(device)teacher_output = teacher_model(data)student_output = student_model(data)soft_loss = KD_loss(teacher_output, student_output)hard_loss = HL(student_output, labels)loss = hard_loss + alpha*soft_lossoptimizer_student.zero_grad()loss.backward()optimizer_student.step()print("KD loss: ", loss)student_model.eval()ACC = 0for data, labels in test_loader:with torch.no_grad():data = data.to(device)labels = labels.to(device)output = student_model(data)_, pred = torch.max(output, dim=1)ACC += float(torch.sum(pred == labels))print('KD test_acc:{}'.format(ACC / len(test_dataset)))

代码中的teacher_output是teacher网络的输出，student_output是student的输出，两者之间设计的KD_loss代码如下：

KD_loss代码：

Temp为温度系数，默认为2【可以根据自己的数据集去尝试】

alpha是hard与soft的平衡系数【默认0.5，也是根据自己的实际情况调整】

损失函数采用的KL，你也可以改为交叉熵。

Temp = 2.  # 温度常数
alpha = 0.5
def KD_loss(p, q):  # p指的老师老师的预测(经过softmax)，q是学生的预测pt = F.softmax(p / Temp, dim=1)ps = F.log_softmax(q / Temp, dim=1)return nn.KLDivLoss(reduction='mean')(ps, pt) * (Temp**2)

KD tran后student结果： 

KD loss:  tensor(0.2580, device='cuda:0', grad_fn=)
KD test_acc:0.9753
KD loss:  tensor(0.1686, device='cuda:0', grad_fn=)
KD test_acc:0.9748
KD loss:  tensor(0.0827, device='cuda:0', grad_fn=)
KD test_acc:0.9849
KD loss:  tensor(0.0098, device='cuda:0', grad_fn=)
KD test_acc:0.9865
KD loss:  tensor(0.0114, device='cuda:0', grad_fn=)
KD test_acc:0.988

 可以看出经过KD训练后student略有提升【主要手写数字这个太容易训练，稍微一训练就可以有较高的准确率】，如果换成别的数据集【比如猫狗数据集可能会明显点，可以自己试试】。

如果要换teacher和student网络，只需要在代码中将teacher_model和student_model网络进行替换即可。

完整代码

目标检测方面的KD比较麻烦，这个以后再讲。

import torchfrom torch.optim import Adam, SGD
import torch.nn.functional as F
import torch.nn as nn
from torchvision.models import resnet50, resnet34, resnet18, MobileNetV2
import torchvision
import torchvision.transforms as transformsTemp = 2.  # 温度常数
alpha = 0.5
def KD_loss(p, q):  # p指的老师老师的预测(经过softmax)，q是学生的预测pt = F.softmax(p / Temp, dim=1)ps = F.log_softmax(q / Temp, dim=1)return nn.KLDivLoss(reduction='mean')(ps, pt) * (Temp**2)
def teacher_train(teacher_model, train_loader, test_loader, loss_func, epochs):teacher_model.train()teacher_model.cuda()# trainfor i in range(epochs):for data, label in train_loader:data = data.to(device)label = label.to(device)output = teacher_model(data)loss = loss_func(output, label)optimizer_teacher.zero_grad()loss.backward()optimizer_teacher.step()print("loss: ", loss)# evalcorrect = 0teacher_model.eval()teacher_model.cuda()for test_data, test_label in test_loader:test_data = test_data.to(device)test_label = test_label.to(device)with torch.no_grad():output = teacher_model(test_data)# acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))# print("teacher acc: ", acc)_, pred = torch.max(output, dim=1)correct += float(torch.sum(pred == test_label))print('test_acc:{}'.format(correct / len(test_dataset)))return teacher_modeldef student_train(student_model, train_loader, test_loader, loss_func, epochs):student_model.train()student_model.cuda()# trainfor i in range(epochs):for data, label in train_loader:data = data.to(device)label = label.to(device)output = student_model(data)loss = loss_func(output, label)optimizer_student.zero_grad()loss.backward()optimizer_student.step()print("student loss: ", loss)# evalcorrect = 0student_model.eval()student_model.cuda()for test_data, test_label in test_loader:test_data = test_data.to(device)test_label = test_label.to(device)with torch.no_grad():output = student_model(test_data)# acc = torch.mean((torch.argmax(F.softmax(output, dim=-1), dim=-1) == test_label).type(torch.FloatTensor))# print("teacher acc: ", acc)_, pred = torch.max(output, dim=1)correct += float(torch.sum(pred == test_label))print('student test_acc:{}'.format(correct / len(test_dataset)))def KD_train(teacher_model, student_model, train_loader, test_loader,loss_func, epochs):teacher_model.eval()student_model.train()student_model.cuda()HL = nn.CrossEntropyLoss()for i in range(epochs):for data, labels in train_loader:data = data.to(device)labels = labels.to(device)teacher_output = teacher_model(data)student_output = student_model(data)soft_loss = KD_loss(teacher_output, student_output)hard_loss = HL(student_output, labels)loss = hard_loss + alpha*soft_lossoptimizer_student.zero_grad()loss.backward()optimizer_student.step()print("KD loss: ", loss)student_model.eval()ACC = 0for data, labels in test_loader:with torch.no_grad():data = data.to(device)labels = labels.to(device)output = student_model(data)_, pred = torch.max(output, dim=1)ACC += float(torch.sum(pred == labels))print('KD test_acc:{}'.format(ACC / len(test_dataset)))def do_train(teacher_model, student_model, train_loader, test_loader, loss_func, epochs):#教师训练teacher_model.train()teacher_model.to(device)print("teacher model train")Teacher = teacher_train(teacher_model, train_loader, test_loader, loss_func, epochs)print("teacher model trained finished!")# print("\n student model ready train")# student_train(student_model, train_loader, test_loader, loss_func, epochs)# print("\n student model trained finished!")print("\n KD model ready train")KD_train(Teacher, student_model, train_loader, test_loader, loss_func, epochs)if __name__=="__main__":# 准备数据集batch_size = 64train_dataset = torchvision.datasets.MNIST('./data/', train=True, download=True,transform=transforms.Compose([transforms.Resize(28),transforms.ToTensor(),transforms.Lambda(lambda x: x.repeat(3, 1, 1)),transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),transforms.Grayscale(num_output_channels=3)]))test_dataset = torchvision.datasets.MNIST('./data/', train=False, download=True,transform=transforms.Compose([transforms.Resize(28),  # resnet默认图片输入大小224*224transforms.ToTensor(),transforms.Lambda(lambda x: x.repeat(3, 1, 1)),transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),transforms.Grayscale(num_output_channels=3)]))train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=batch_size, shuffle=True)test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=batch_size, shuffle=False)sample, label = next(iter(train_loader))print(sample.shape)print("当前类： ", label)num_classes = 10lr = 0.01epochs = 5device = torch.device('cuda:0')teacher_model = resnet18(num_classes=num_classes)student_model = resnet50(num_classes=num_classes)optimizer_teacher = SGD(teacher_model.parameters(), lr=lr, momentum=0.9)optimizer_student = SGD(student_model.parameters(), lr=lr, momentum=0.9)loss_function = nn.CrossEntropyLoss()do_train(teacher_model, student_model, train_loader, test_loader, loss_function, epochs)