YOLOv5:对yolov5n模型进一步剪枝压缩

YOLOv5:对yolov5n模型进一步剪枝压缩

前言

前提条件

相关介绍

  • Python是一种跨平台的计算机程序设计语言。是一个高层次的结合了解释性、编译性、互动性和面向对象的脚本语言。最初被设计用于编写自动化脚本(shell),随着版本的不断更新和语言新功能的添加,越多被用于独立的、大型项目的开发。
  • PyTorch 是一个深度学习框架,封装好了很多网络和深度学习相关的工具方便我们调用,而不用我们一个个去单独写了。它分为 CPU 和 GPU 版本,其他框架还有 TensorFlow、Caffe 等。PyTorch 是由 Facebook 人工智能研究院(FAIR)基于 Torch 推出的,它是一个基于 Python 的可续计算包,提供两个高级功能:1、具有强大的 GPU 加速的张量计算(如 NumPy);2、构建深度神经网络时的自动微分机制。
  • YOLOv5是一种单阶段目标检测算法,该算法在YOLOv4的基础上添加了一些新的改进思路,使其速度与精度都得到了极大的性能提升。它是一个在COCO数据集上预训练的物体检测架构和模型系列,代表了Ultralytics对未来视觉AI方法的开源研究,其中包含了经过数千小时的研究和开发而形成的经验教训和最佳实践。
  • 剪枝是一种通过去除网络中冗余的channels,filters, neurons, or layers以得到一个更轻量级的网络,同时不影响性能的方法。

具体步骤

修改yolov5n.yaml配置文件

yaml 复制代码
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license

# Parameters
nc: 80  # number of classes
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.50  # layer channel multiple
bash 复制代码
 - nc: 8 代表数据集中的类别数目。
 - depth_multiple: 0.33
 	- 用来控制模型的深度,仅在number≠1时启用。
 	- 如第一个C3层的参数设置为[-1, 3, C3, [128]],其中number=3,表示在yolov5s中含有 3 × 0.33 ≈ 1个C3。
 - width_multiple: 0.50
 	- 用来控制模型的宽度,主要作用于args中的channel_out。
 	- 如第一个Conv层,输出通道数channel_out=64,那么在yolov5s中,会将卷积过程中的卷积核设置为 64 × 0.50 = 32,所以会输出 32 通道的特征图。

yaml 复制代码
depth_multiple: 0.33  # model depth multiple
width_multiple: 0.25  # layer channel multiple

改为

yaml 复制代码
depth_multiple: 0.16  # model depth multiple
width_multiple: 0.125  # layer channel multiple

即可达到减少卷积层数 的目的。

单通道数据(黑白图片)

  • 如果数据集是单通道数据,即黑白图片数据集,还可以修改训练时输入的通道数(yolov5默认输入通道数ch=3,我们可以修改ch=1),减少训练参数。
  • 如果是彩色图片数据集,可跳过此部分的内容

修改models/yolo.py文件

python 复制代码
        if m in {
                Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
                BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x}:
            c1, c2 = ch[f], args[0]
            if c2 != no:  # if not output
                c2 = make_divisible(c2 * gw, 8)

添加:

python 复制代码
        if m in (Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv,
                 BottleneckCSP, C3, C3TR, C3SPP, C3Ghost, nn.ConvTranspose2d, DWConvTranspose2d, C3x):
            c1, c2 = ch[f], args[0]
            # 添加的内容
            if i == 0: # 第一层输入,为单通道图片
                c1 = 1
            if c2 != no:  # if not output
                c2 = make_divisible(c2 * gw, 8)
python 复制代码
        # Define model
        ch = self.yaml['ch'] = self.yaml.get('ch', ch)  # input channels

添加:

python 复制代码
        # Define model
        ch = self.yaml['ch'] = self.yaml.get('ch', ch)  # input channels

        # 添加的内容
        self.yaml['ch'] = 1
        ch = self.yaml['ch']

修改train.py文件

python 复制代码
model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create

改为

python 复制代码
# 修改的内容
# model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
model = Model(cfg or ckpt['model'].yaml, ch=1, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
python 复制代码
model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create

改为

python 复制代码
# 修改的内容
# model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
model = Model(cfg, ch=1, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
python 复制代码
for i, (imgs, targets, paths, _) in pbar:  # batch -------------------------------------------------------------
            callbacks.run('on_train_batch_start')
            ni = i + nb * epoch  # number integrated batches (since train start)
            imgs = imgs.to(device, non_blocking=True).float() / 255  # uint8 to float32, 0-255 to 0.0-1.0

            # Warmup
            if ni <= nw:
                xi = [0, nw]  # x interp
                # compute_loss.gr = np.interp(ni, xi, [0.0, 1.0])  # iou loss ratio (obj_loss = 1.0 or iou)
                accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
                for j, x in enumerate(optimizer.param_groups):
                    # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
                    x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
                    if 'momentum' in x:
                        x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])

添加:

python 复制代码
for i, (imgs, targets, paths, _) in pbar:  # batch -------------------------------------------------------------
            callbacks.run('on_train_batch_start')
            ni = i + nb * epoch  # number integrated batches (since train start)
            # imgs = imgs.to(device, non_blocking=True).float() / 255  # uint8 to float32, 0-255 to 0.0-1.0

            # 添加的内容,目的是将训练集的图片变为单通道图片(黑白图片)
            imgs = imgs[:, 0, :, :].unsqueeze(1).to(device, non_blocking=True).float() / 255 

            # Warmup
            if ni <= nw:
                xi = [0, nw]  # x interp
                # compute_loss.gr = np.interp(ni, xi, [0.0, 1.0])  # iou loss ratio (obj_loss = 1.0 or iou)
                accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
                for j, x in enumerate(optimizer.param_groups):
                    # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
                    x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
                    if 'momentum' in x:
                        x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])

剪枝后模型大小

  • 原来的yolo5n模型大小为3.5m,剪枝训练后的yolo5n模型大小为2.6m。

参考

[1] https://github.com/ultralytics/yolov5

相关推荐
_.Switch14 分钟前
高级Python自动化运维:容器安全与网络策略的深度解析
运维·网络·python·安全·自动化·devops
AI极客菌44 分钟前
Controlnet作者新作IC-light V2:基于FLUX训练,支持处理风格化图像,细节远高于SD1.5。
人工智能·计算机视觉·ai作画·stable diffusion·aigc·flux·人工智能作画
阿_旭1 小时前
一文读懂| 自注意力与交叉注意力机制在计算机视觉中作用与基本原理
人工智能·深度学习·计算机视觉·cross-attention·self-attention
王哈哈^_^1 小时前
【数据集】【YOLO】【目标检测】交通事故识别数据集 8939 张,YOLO道路事故目标检测实战训练教程!
前端·人工智能·深度学习·yolo·目标检测·计算机视觉·pyqt
测开小菜鸟1 小时前
使用python向钉钉群聊发送消息
java·python·钉钉
Power20246662 小时前
NLP论文速读|LongReward:基于AI反馈来提升长上下文大语言模型
人工智能·深度学习·机器学习·自然语言处理·nlp
数据猎手小k2 小时前
AIDOVECL数据集:包含超过15000张AI生成的车辆图像数据集,目的解决旨在解决眼水平分类和定位问题。
人工智能·分类·数据挖掘
好奇龙猫2 小时前
【学习AI-相关路程-mnist手写数字分类-win-硬件:windows-自我学习AI-实验步骤-全连接神经网络(BPnetwork)-操作流程(3) 】
人工智能·算法
沉下心来学鲁班2 小时前
复现LLM:带你从零认识语言模型
人工智能·语言模型
数据猎手小k2 小时前
AndroidLab:一个系统化的Android代理框架,包含操作环境和可复现的基准测试,支持大型语言模型和多模态模型。
android·人工智能·机器学习·语言模型