我们在训练好了模型之后，在模型部署上线之前，我们通常会对模型进行优化。例如：我们的模型过大，使得模型无法部署到边缘设备上，我们需要对模型进行压缩，那么怎么对模型进行压缩和优化呢？
我们需要在减小模型大小的同时，尽可能的维持模型原有精度。
模型优化的方法：
- 模型剪枝：深度学习模型压缩之模型剪枝 - 简书 (jianshu.com)
- 模型量化：深度学习模型压缩之模型量化 - 简书 (jianshu.com)
- 模型蒸馏：深度学习模型压缩之模型蒸馏 - 简书 (jianshu.com)
一、知识蒸馏
文章中的第一张图片描述的就是模型蒸馏的一个过程。
我们先用一个大模型去拟合一个任务，将精度提升。然后将这个大模型作为老师网络，设计一个小模型作为学生网络。将老师模型的输出作为标签，让学生去学习。这就是知识蒸馏。
从复杂网络（teacher net）中抽取徐连数据的分布交给简单网络（student net）。
我们在训练老师模型的时候我们使用onehot做标签，onehot被称之为硬标签（hard targe）其实这个标签并不是很好，因为他无法表示每个标签之间的联系。然而我们训练好的大模型的输出就是一个软目标(soft target)，因为它能够极好的表示每一个物品分类之间的概率。为了能够使得软目标更加软，引入蒸馏温度T。
当T=1时，公式就等于了Softmax。
可参考：同济子豪兄的视频 【精读AI论文】知识蒸馏_哔哩哔哩_bilibili
下图就是知识蒸馏的全过程：
首先，上面是教师网络，这个网络是提前训练好的，下面的是学生网络。我们将同样的数据输入给教师网络和学生网络，在蒸馏温度T下，得出两个结果soft labels 和soft predictions。这两者之间求损失进行梯度更新。同时学生网络需要与真实标签再求一次损失，进行梯度更新。这样下来，学生网络一边在学习老师网络总结出来的规律，一边在真实数据中进行训练，能够较好地提取老师网络中有用的权重参数。
第一个损失函数，我们使用的是KL散度（即，相对熵）来衡量两个分布之间的差别。因为交叉熵损失函数仅描述了两个分布是否相似，但是并不描述两个分布的位置是否相同。第二个损失函数，我们使用的是普通的交叉熵损失函数。
示例代码：
from typing import Any
import torch
from torch import nn, optim
from torch.nn import CrossEntropyLoss, KLDivLoss
from torch.functional import F
from torch.utils.data import DataLoader, TensorDataset, SequentialSampler
# 学生网络
class Net_V1(nn.Module):
def __init__(self) -> None:
super().__init__()
self.layers = nn.Sequential(
nn.Linear(784, 512), nn.ReLU(), nn.BatchNorm1d(512),
nn.Linear(512, 256), nn.ReLU(), nn.BatchNorm1d(256),
nn.Linear(256, 10)
)
def forward(self, x):
return self.layers(x)
# 教师网络
class Net_V2(nn.Module):
def __init__(self) -> None:
super().__init__()
self.layers = nn.Sequential(
nn.Linear(784, 512), nn.ReLU(), nn.BatchNorm1d(512),
nn.Linear(512, 256), nn.ReLU(), nn.BatchNorm1d(256),
nn.Linear(256, 128), nn.ReLU(), nn.BatchNorm1d(128),
nn.Linear(128, 64), nn.ReLU(), nn.BatchNorm1d(64),
nn.Linear(64, 32), nn.ReLU(), nn.BatchNorm1d(32),
nn.Linear(32, 10)
)
def forward(self, x):
return self.layers(x)
# 实例化模型
student_model = Net_V1()
teacher_model = Net_V2() # 这里假装它是训练好的模型
torch.manual_seed(13)
# 构建输入数据和标签
inputs = torch.randn(8,28*28)
true_label = torch.randint(0,10,size=(8,))
dataset = TensorDataset(inputs, true_label)
sampler = SequentialSampler(inputs)
dataloader = DataLoader(dataset, sampler=sampler, batch_size=2)
loss_fn_ce = CrossEntropyLoss()
loss_fn_kl = KLDivLoss(reduction="mean")
opt = optim.Adam(student_model.parameters())
a = 0.9
t=7 # 蒸馏温度
for epoch in range(20):
for i,(x,y) in enumerate(dataloader):
out_student = student_model(x)
out_teacher = teacher_model(x)
loss_hard = loss_fn_ce(out_student,y)
loss_soft = loss_fn_kl(F.log_softmax(out_student/t, dim=1), F.log_softmax(out_teacher/t, dim=1))
loss = a*loss_soft + (1-a)*loss_hard
opt.zero_grad()
loss.backward()
opt.step()
print(f"epoch: {epoch}, batch: {i}, loss: {loss.item()}")
