首先蒸馏是什么？
模型蒸馏（Model Distillation）是一种用于在计算机视觉中提高模型性能和效率的技术。在模型蒸馏中，通常存在两个模型，即“教师模型”和“学生模型”。
为什么需要蒸馏？
- 在不增加模型计算量和参数量的情况下提升精度，也即是可以无损提高精度。
- 配合剪枝一起使用，可以尽量达到无损降低模型参数量、计算量，提高FPS的情况下，还能保持模型精度没有下降甚至上升，这是改进网络结构无法达到的高度。
- 论文中的保底手段，因为剪枝和蒸馏的特殊性，其都不会增加参数量和计算量，可以在最后一个点上大幅度增加实验和工作量，因为本身蒸馏也需要做大量实验。
目录
（1）本文选取的老师模型为yolov8s，学生为剪枝完的yolov8s
（3）使用前下载必须的包，并且把数据集放在datasets文件夹中，最后替换data.yaml中分类。
（2）约束训练后，先进行剪枝，使用prune.py。替换模型位置，直接运行。
（3）剪完枝后，效果不一定好，所以使用剪枝完后的模型，继续训练：
一.代码前提
（1）本文选取的老师模型为yolov8s，学生为剪枝完的yolov8s
（2）本文使用的蒸馏方法包括mgd，cwd
（3）使用前下载必须的包，并且把数据集放在datasets文件夹中，最后替换data.yaml中分类。
本文代码已经上传到GitHub，链接：yolov8_蒸馏
使用不妨加个关注，后续还会加入Vit（vision transformer），替换loss等提升精度的方法。
二.蒸馏步骤
(1) 训练教师模型
打开文件中train.py，替换模型文件位置。开始训练，达到理想目标就停止。
import os
from ultralytics import YOLO
import torch
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
def main():
model = YOLO("yolov8s.pt")
model.train(data="data.yaml", Distillation = None, loss_type='None', amp=False, imgsz=640, epochs=50, batch=20, device=0, workers=0)
if __name__ == '__main__':
main()
(2) 训练学生模型
打开文件中train.py，替换模型文件位置。我这边使用的是剪枝后的yolov8s模型，具体轻量化剪枝步骤可见本文最后。
import os
from ultralytics import YOLO
import torch
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
def main():
model_s = YOLO("./runs/detect/prune/weights/prune.pt")
model_s.train(data="data.yaml", Distillation = None, loss_type='None', amp=False, imgsz=640, epochs=50, batch=20, device=0, workers=0)
if __name__ == '__main__':
main()
(3) 蒸馏训练
打开文件中train_distillation.py，替换老师与学生模型文件位置。两种蒸馏方法可以选择：cwd和mgd。
import os
from ultralytics import YOLO
import torch
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
def main():
model_t = YOLO('runs/detect/yolov8s/weights/best.pt') # the teacher model
model_s = YOLO('runs/detect/prune/weights/best.pt') # the student model
"""
Attributes:
Distillation: the distillation model
loss_type: mgd, cwd
amp: Automatic Mixed Precision
"""
model_s.train(data="data.yaml", Distillation=model_t.model, loss_type='mgd', amp=False, imgsz=640, epochs=100,
batch=20, device=0, workers=0, lr0=0.001)
if __name__ == '__main__':
main()
现在先不进行训练，打开文件夹yolo_project_distillation\ultralytics\engine\trainer.py
在类FeatureLoss中，函数forward大概162行处打一个断点，进行调试。代码位置：
def forward(self, y_s, y_t):
assert len(y_s) == len(y_t)
tea_feats = []
stu_feats = []
for idx, (s, t) in enumerate(zip(y_s, y_t)):
# change ---
if self.distiller == 'cwd':
s = self.align_module[idx](s)
s = self.norm[idx](s)
else:
s = self.norm1[idx](s)
t = self.norm[idx](t)
tea_feats.append(t)
stu_feats.append(s)
loss = self.feature_loss(stu_feats, tea_feats)
return self.loss_weight * loss调试运行，查看变量中学生模型y_s和老师模型y_t的张量大小。把通道数记下来，写在类Distillation_loss的
channels_s = [256, 480, 256, 64, 143, 229][-le:]
channels_t = [256, 512, 256, 128, 256, 512][-le:]这边总共有六个，刚好对应模型的六个层的通道数。
替换完成后，应该就可以进行训练了。训练不好的话，再来评论区找我吧。
三.模型剪枝+蒸馏
（1）约束训练在我上一篇文章中提到，链接：yolov8剪枝
（2）约束训练后，先进行剪枝，使用prune.py。替换模型位置，直接运行。
from ultralytics import YOLO
import torch
from ultralytics.nn.modules import Bottleneck, Conv, C2f, SPPF, Detect
from copy import deepcopy
# Load a model
yolo = YOLO("./runs/detect/yolov8s/weights/last.pt")
# Save model address
res_dir = "./runs/detect/prune/weights/prune.pt"
# Pruning rate
factor = 0.75
yolo.info()
model = yolo.model
ws = []
bs = []
for name, m in model.named_modules():
if isinstance(m, torch.nn.BatchNorm2d):
w = m.weight.abs().detach()
b = m.bias.abs().detach()
ws.append(w)
bs.append(b)
# print(name, w.max().item(), w.min().item(), b.max().item(), b.min().item())
# keep
ws = torch.cat(ws)
threshold = torch.sort(ws, descending=True)[0][int(len(ws) * factor)]
print(threshold)
def prune_conv(conv1: Conv, conv2: Conv):
gamma = conv1.bn.weight.data.detach()
beta = conv1.bn.bias.data.detach()
keep_idxs = []
local_threshold = threshold
while len(keep_idxs) < 8:
keep_idxs = torch.where(gamma.abs() >= local_threshold)[0]
local_threshold = local_threshold * 0.5
n = len(keep_idxs)
# n = max(int(len(idxs) * 0.8), p)
# print(n / len(gamma) * 100)
# scale = len(idxs) / n
conv1.bn.weight.data = gamma[keep_idxs]
conv1.bn.bias.data = beta[keep_idxs]
conv1.bn.running_var.data = conv1.bn.running_var.data[keep_idxs]
conv1.bn.running_mean.data = conv1.bn.running_mean.data[keep_idxs]
conv1.bn.num_features = n
conv1.conv.weight.data = conv1.conv.weight.data[keep_idxs]
conv1.conv.out_channels = n
if conv1.conv.bias is not None:
conv1.conv.bias.data = conv1.conv.bias.data[keep_idxs]
if not isinstance(conv2, list):
conv2 = [conv2]
for item in conv2:
if item is not None:
if isinstance(item, Conv):
conv = item.conv
else:
conv = item
conv.in_channels = n
conv.weight.data = conv.weight.data[:, keep_idxs]
def prune(m1, m2):
if isinstance(m1, C2f): # C2f as a top conv
m1 = m1.cv2
if not isinstance(m2, list): # m2 is just one module
m2 = [m2]
for i, item in enumerate(m2):
if isinstance(item, C2f) or isinstance(item, SPPF):
m2[i] = item.cv1
prune_conv(m1, m2)
for name, m in model.named_modules():
if isinstance(m, Bottleneck):
prune_conv(m.cv1, m.cv2)
seq = model.model
for i in range(3, 9):
if i in [6, 4, 9]: continue
prune(seq[i], seq[i + 1])
detect: Detect = seq[-1]
last_inputs = [seq[15], seq[18], seq[21]]
colasts = [seq[16], seq[19], None]
for last_input, colast, cv2, cv3 in zip(last_inputs, colasts, detect.cv2, detect.cv3):
prune(last_input, [colast, cv2[0], cv3[0]])
prune(cv2[0], cv2[1])
prune(cv2[1], cv2[2])
prune(cv3[0], cv3[1])
prune(cv3[1], cv3[2])
for name, p in yolo.model.named_parameters():
p.requires_grad = True
#yolo.val(workers=0) # 剪枝模型进行验证 yolo.val(workers=0)
yolo.info()
# yolo.export(format="onnx") # 导出为onnx文件
# yolo.train(data="./data/data_nc5/data_nc5.yaml", epochs=100) # 剪枝后直接训练微调
ckpt = {
'epoch': -1,
'best_fitness': None,
'model': yolo.ckpt['ema'],
'ema': None,
'updates': None,
'optimizer': None,
'train_args': yolo.ckpt["train_args"], # save as dict
'date': None,
'version': '8.0.142'}
torch.save(yolo.ckpt, res_dir)
（3）剪完枝后，效果不一定好，所以使用剪枝完后的模型，继续训练：
import os
from ultralytics import YOLO
import torch
os.environ['KMP_DUPLICATE_LIB_OK'] = 'True'
def main():
# model = YOLO(r'ultralytics/cfg/models/v8/yolov8s.yaml').load('runs/detect/yolov8s/weights/best.pt')
model_s = YOLO("./runs/detect/prune/weights/prune.pt")
model_s.train(data="data.yaml", Distillation = None, loss_type='None', amp=False, imgsz=640, epochs=50, batch=20, device=0, workers=0)
if __name__ == '__main__':
main()
------------------------------------------over！！！！！！！！！！！！！！！！！------------------------------
710
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
