知识蒸馏——pytorch实现
轻量化网络
知识蒸馏可以理解为轻量化网络的一个tricks，轻量化网络是深度学习的一个大的发展趋势，尤其是在移动端，终端边缘计算这种对算力和运算时间有要求的场景中。
轻量化网络可以有以下四种方式实现：
1.压缩已经训练好的模型: 知识蒸馏，权值量化，剪枝，注意力迁移
2.直接训练轻量化网络：SqueezeNet，MobileNet等
3.加速卷积运算：低秩分解
4.硬件部署：Tensorrt，Jetson，Openvino等
知识蒸馏
知识蒸馏在轻量化网络方面具有很高的地位。下图是知识蒸馏的实现过程。
导入包
import torch
from torch import nn
import torch.nn.functional as F
import torchvision
from torchvision import transforms
from torch.utils.data import DataLoader
from torchinfo import summary
from tqdm import tqdm
#设置随机数种子，便于复现
torch.manual_seed(0)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# 使用cudnn加速卷积运算
torch.backends.cudnn.benchmark = True
载入MNIST数据集
from torchvision.transforms.transforms import ToTensor
# 载入训练集
train_dataset = torchvision.datasets.MNIST(
root = 'dataset/',
train = True,
transform = transforms.ToTensor(),
download=True
)
# 生成测试集
test_dataset = torchvision.datasets.MNIST(
root = 'dataset/',
train = False,
transform = transforms.ToTensor(),
download=True
)
# 生成dataloader
train_dataloader = DataLoader(dataset=train_dataset,batch_size=32,shuffle=True)
test_dataloader = DataLoader(dataset=test_dataset,batch_size=32,shuffle=True)
构建教师模型
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)
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 x
训练教师模型
model = TeacherModel()
model = model.to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr=1e-4)
epochs = 6
for epoch in range(epochs):
model.train()
# 在训练集上训练
for data, targets in tqdm(train_dataloader):
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 = 0
num_samples = 0
with 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()
model.train()
print('Epoch:{}\t Accuracy:{:.4f}'.format(epoch+1,acc))
教师模型预测结果
创建学生模型
class 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)
self.dropout = nn.Dropout(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 x
训练学生模型
model = StudentModel()
model = model.to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(),lr=1e-4)
epochs = 6
for epoch in range(epochs):
model.train()
# 在训练集上训练
for data, targets in tqdm(train_dataloader):
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 = 0
num_samples = 0
with 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()
model.train()
print('Epoch:{}\t Accuracy:{:.4f}'.format(epoch+1,acc))
学生模型预测结果
学生模型由于比教师模型更加轻量化（教师模型的隐藏层有1200个神经元，学生模型的只有20个神经元），所以性能并不如教师模型
student_model_scratch = model
知识蒸馏训练模型
# 准备预训练好的教师模型
teacher_model.eval()
# 准备新的学生模型
model = StudentModel()
model = model.to(device)
model.train()
# 蒸馏温度
temp = 7
# hard_loss
hard_loss = nn.CrossEntropyLoss()
# hard_loss权重
alpha = 0.3
# soft_loss
soft_loss = nn.KLDivLoss(reduction='batchmean')
optimizer = torch.optim.Adam(model.parameters(),lr=1e-4)
epochs = 10
for epoch in range(epochs):
# 训练集上训练模型权重
for data,targets in tqdm(train_dataloader):
data = data.to(device)
targets = targets.to(device)
# 教师模型预测
with torch.no_grad():
teachers_preds = teacher_model(data)
# 学生模型预测
students_preds = model(data)
# 计算hard_loss
students_loss = hard_loss(students_preds,targets)
# 计算蒸馏后的预测结果及soft_loss
ditillation_loss = soft_loss(
F.softmax(students_preds/temp,dim=1),
F.softmax(teachers_preds/temp,dim=1)
)
# 将hard_loss和soft_loss加权求和
loss = alpha*students_loss+(1-alpha)*ditillation_loss
# 反向传播，优化权重
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 测试集上评估模型性能
model.eval()
num_correct = 0
num_samples = 0
with 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()
model.train()
print('Epoch:{}\t Accuracy:{:.4f}'.format(epoch+1,acc))
知识蒸馏训练后的预测结果
虽然结果差不多，但是这只是知识蒸馏的一个小的应用，还有就是MNIST数据并不是很多，所以导致差异不明显，但是可以通过这个代码更好的了解知识蒸馏模型。知识蒸馏绝对是轻量化网络的挖坑之作。