yolov8蒸馏（附代码-免费）

数学人学python

已于 2024-08-04 10:59:44 修改

阅读量9.2k

点赞数 59

CC 4.0 BY-SA版权

文章标签： YOLO python 算法

于 2024-08-04 10:57:46 首次发布

本文链接：https://blog.youkuaiyun.com/W_extend/article/details/140902235

首先蒸馏是什么？

模型蒸馏（Model Distillation）是一种用于在计算机视觉中提高模型性能和效率的技术。在模型蒸馏中，通常存在两个模型，即“教师模型”和“学生模型”。

为什么需要蒸馏？

在不增加模型计算量和参数量的情况下提升精度，也即是可以无损提高精度。
配合剪枝一起使用，可以尽量达到无损降低模型参数量、计算量，提高FPS的情况下，还能保持模型精度没有下降甚至上升，这是改进网络结构无法达到的高度。
论文中的保底手段，因为剪枝和蒸馏的特殊性，其都不会增加参数量和计算量，可以在最后一个点上大幅度增加实验和工作量，因为本身蒸馏也需要做大量实验。

一.代码前提

（1）本文选取的老师模型为yolov8s，学生为剪枝完的yolov8s

（2）本文使用的蒸馏方法包括mgd，cwd

（3）使用前下载必须的包，并且把数据集放在datasets文件夹中，最后替换data.yaml中分类。

（1）约束训练在我上一篇文章中提到，链接：yolov8剪枝

（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！！！！！！！！！！！！！！！！！------------------------------