YOLOX改进之模型轻量化(Lite)

文章内容:给YOLOv5-Lite系列轻量级模型换上YOLOX头部
环境:pytorch1.7+cuda11.0
注意:YOLOv5版本为5.0,可以匹配YOLOX的主干网络与neck部分

一、参考资料

参考资料1

YOLOv5-Lite
代码链接https://github.com/ppogg/YOLOv5-Lite
文章链接https://blog.csdn.net/weixin_45829462/article/details/119767896?spm=1001.2014.3001.5502

参考资料2
YOLOv5-YOLOX(yolov5代码风格复现版):
代码链接https://gitee.com/SearchSource/yolov5_yolox
文章链接https://www.yuque.com/yerunyuan/ar9831/tsm0id#Kfi4w

参考资料3
YOLOX原理解析(大白老师版):https://zhuanlan.zhihu.com/p/397993315

二、代码修改位置

运行代码:YOLOv5-YOLOX版本(把YOLOv5-Lite改进代码移过来更方便)
YOLOv5-Lite:(先让改进的Lite代码在YOLOv5-YOLOX版本上运行起来)
1、移动YOLOv5-Lite中模型结构改进代码(位于common.py中)
2、修改yolo.py代码,添加新的yaml模型文件读取模块(如:RepVGG等,位于300多行),可以参考YOLOv5代码修改。
YOLOv5-Lite各系列性能
(图来自https://github.com/ppogg/YOLOv5-Lite)
在这里插入图片描述

YOLOX_s-Lite
(在改进的Lite模型上更换YOLOX-s头部,包括c、e、g、s)
很简单,直接添加和简略修改,仔细观察yolox.yaml文件格式,对应修改,yoloxs的头部通道为128;其次保证yolox.py代码中通道为128(33行左右)。yaml文件修改如下图:
在这里插入图片描述

YOLOX_nano-Lite
(在改进的Lite模型上更换YOLOX-s头部,包括e、s)
修改位置同上,yaml文件的head通道均改成64;以及yolox.py代码(33行左右)head通道修改为64,如下图:
在这里插入图片描述

在这里插入图片描述

三、运行代码命令

# 命令:
# 一、v5系列改进
# v5Lite-c:(8.66M)
python train.py  --img-size 640  --data PCB.yaml --cfg models/v5Lite-c.yaml --hyp data/hyps/hyp.scratch.yaml --weights v5lite-c.pt --batch-size 8 --epochs 1 --device 0

# v5Lite-e:(1.62M)
python train.py  --img-size 640  --data PCB.yaml --cfg models/v5Lite-e.yaml --hyp data/hyps/hyp.scratch.yaml --weights v5lite-e.pt --batch-size 8 --epochs 1 --device 0

# v5Lite-g:(10.7M)
python train.py  --img-size 640  --data PCB.yaml --cfg models/v5Lite-g.yaml --hyp data/hyps/hyp.scratch.yaml --weights v5lite-g.pt --batch-size 8 --epochs 1 --device 0

# v5Lite-s:(3.25M)
python train.py  --img-size 640  --data PCB.yaml --cfg models/v5Lite-s.yaml --hyp data/hyps/hyp.scratch.yaml --weights v5lite-s.pt --batch-size 8 --epochs 1 --device 0


# 二、yolox及改进系列:

# yoloxs-office(纯silu,数据增强、精度与官网接近)(17.2M)
python train.py --noautoanchor --img-size 640  --data PCB.yaml --cfg models/yoloxs_official.yaml --hyp data/hyps/hyp.scratch.yolox.official.yaml --weights yolox-s.pt --batch-size 8 --epochs 1 --device 0

# yoloxs_rslu(激活函数:relu+silu,且数据增强有所更改,权重不变,推理速度加快,由于激活函数改变,精度稍微降低)(17.2M)
python train.py --noautoanchor --img-size 640  --data PCB.yaml --cfg models/yoloxs_rslu.yaml --hyp data/hyps/hyp.scratch.yolox.yaml --weights yolox-s.pt --batch-size 8 --epochs 1 --device 0

# yoloxs(激活函数:纯silu,且数据增强有所更改,权重不变,修改后训练速度加快,精度与官网接近)(17.2M)
python train.py --noautoanchor --img-size 640  --data PCB.yaml --cfg models/yoloxs.yaml --hyp data/hyps/hyp.scratch.yolox.yaml --weights yolox-s_rslu.pt --batch-size 8 --epochs 1 --device 0


# 【改进1】:前面v5的改进换上了yoloxs(头部通道为128)的head,anchor_free锚框机制

# yoloxs_Lite-c:(12.1M)
python train.py --noautoanchor --img-size 640  --data PCB.yaml --cfg models/yoloxs_Lite_c.yaml --hyp data/hyps/hyp.scratch.yolox.yaml --weights v5lite-c.pt --batch-size 8 --epochs 1 --device 0

# yoloxs_Lite-e:(5.04M)
python train.py --noautoanchor --img-size 640  --data PCB.yaml --cfg models/yoloxs_Lite_e.yaml --hyp data/hyps/hyp.scratch.yolox.yaml --weights v5lite-e.pt --batch-size 8 --epochs 1 --device 0

# yoloxs_Lite-g:(14.2M)
python train.py --noautoanchor --img-size 640  --data PCB.yaml --cfg models/yoloxs_Lite_g.yaml --hyp data/hyps/hyp.scratch.yolox.yaml --weights v5lite-g.pt --batch-size 8 --epochs 1 --device 0

# yoloxs_Lite-s:(6.70M)
python train.py --noautoanchor --img-size 640  --data PCB.yaml --cfg models/yoloxs_Lite_s.yaml --hyp data/hyps/hyp.scratch.yolox.yaml --weights v5lite-s.pt --batch-size 8 --epochs 1 --device 0


# 【改进2】:e和s系列换上了yolox-nano(头部通道为64)的head
# yoloxnano_Lite-e:(2.46M)
python train.py --noautoanchor --img-size 640  --data PCB.yaml --cfg models/yoloxnano_Lite_e.yaml --hyp data/hyps/hyp.scratch.yolox.yaml --weights v5lite-e.pt --batch-size 8 --epochs 1 --device 0

# yoloxnano_Lite-s:(4.10M)
python train.py --noautoanchor --img-size 640  --data PCB.yaml --cfg models/yoloxnano_Lite_s.yaml --hyp data/hyps/hyp.scratch.yolox.yaml --weights v5lite-s.pt --batch-size 8 --epochs 1 --device 0

注意】:预训练权重不是lite给的,也能进行训练,上面PCB.yaml就是自己的训练数据集内容。数据集制作与训练可以参考:https://blog.csdn.net/weixin_45679938/article/details/118803745
使用心得】:
1、在我的数据集上,yoloxs-lite-g精度与yolox齐平,高于yolov5有2个点,并且模型大小与yolov5差不多大,模型大小减少了3.2M,其他的没有进行测试,换头后应该都会有提升,毕竟YOLOX相对于v5的改进就在头部;
2、在这个上面添加另外一些trick,如注意力机制、CIOU等,精度有提升(200epoch以上),而相应改进用于官方代码没有提升(原因不明);
3、这边修改模型结构相比官方代码更为简便。

已修改的代码百度网盘链接:(包含v5-Lite系列权重)
链接:https://pan.baidu.com/s/1UfuJWdpkvSvTtjlsDgzLlA
提取码:b718

github链接:https://github.com/Bigtuo/YOLOX-Lite

写的比较急,不喜勿喷,有错误恳求批评指正!

  • 14
    点赞
  • 98
    收藏
    觉得还不错? 一键收藏
  • 39
    评论
你可以在以下代码片段中找到YOLOv8轻量化改进的示例代码: ```python import torch import torch.nn as nn import torch.nn.functional as F class ConvBlock(nn.Module): def __init__(self, in_channels, out_channels, kernel_size, stride, padding): super(ConvBlock, self).__init__() self.conv = nn.Conv2d(in_channels, out_channels, kernel_size, stride, padding) self.bn = nn.BatchNorm2d(out_channels) self.relu = nn.ReLU(inplace=True) def forward(self, x): x = self.conv(x) x = self.bn(x) x = self.relu(x) return x class YOLOv8Lite(nn.Module): def __init__(self, num_classes=80): super(YOLOv8Lite, self).__init__() self.num_classes = num_classes # Backbone self.conv1 = ConvBlock(3, 16, kernel_size=3, stride=1, padding=1) self.conv2 = ConvBlock(16, 32, kernel_size=3, stride=2, padding=1) self.conv3 = ConvBlock(32, 64, kernel_size=3, stride=2, padding=1) self.conv4 = ConvBlock(64, 128, kernel_size=3, stride=2, padding=1) # Head self.conv5 = ConvBlock(128, 256, kernel_size=3, stride=2, padding=1) self.fc1 = nn.Linear(256 * 7 * 7, 4096) self.fc2 = nn.Linear(4096, 7 * 7 * (5 + self.num_classes)) def forward(self, x): x = self.conv1(x) x = self.conv2(x) x = self.conv3(x) x = self.conv4(x) x = self.conv5(x) x = x.view(x.size(0), -1) x = self.fc1(x) x = self.fc2(x) x = x.view(x.size(0), 7, 7, 5 + self.num_classes) return x model = YOLOv8Lite() ``` 这是一个简化版的YOLOv8模型,使用了较少的卷积层和全连接层。你可以根据需要修改网络结构和超参数。请注意,这只是一个示例代码,可能需要根据你的具体任务进行调整。

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论 39
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值