YOLOv5从入门到部署之：训练私有数据并修改模型

最新推荐文章于 2024-09-21 15:18:39 发布

Ring__Rain

最新推荐文章于 2024-09-21 15:18:39 发布

阅读量2k

点赞数 4

分类专栏： pytorch 深度学习

本文链接：https://blog.csdn.net/qq_30263737/article/details/119568674

版权

深度学习同时被 2 个专栏收录

50 篇文章 2 订阅

订阅专栏

pytorch

35 篇文章 6 订阅

订阅专栏

1 环境部署

YOLOv5是基于pytorch实现的, 需要的运行环境已经在requirements.txt文件里面列出。

# base ----------------------------------------
Cython
matplotlib>=3.2.2
numpy>=1.18.5
opencv-python>=4.1.2
pillow
PyYAML>=5.3
scipy>=1.4.1
tensorboard>=2.2
torch>=1.6.0
torchvision>=0.7.0
tqdm>=4.41.0

因此可以直接使用pip进行安装：

pip install -U -r requirements.txt

2 准备数据

YOLOv5的数据集格式和之前的YOLOv3一致。每张图片都有自己的目标的txt文件

例如：

datasets/score/images/train/00b5fefed.jpg     # image
datasets/score/lables/train/00b5fefed.txt       # lable

这里我们使用之前kaggle的小麦头检测的比赛为例，讲解数据准备的流程。

kaggle的小麦头检测的数据集下载链接：https://www.kaggle.com/c/global-wheat-detection

这里值得注意的是，Yolov5的开源许可是GPL-3.0 License不满足Kaggle比赛的要求，因此不能在Kaggle比赛中使用，本文只是为了做个训练示范。

2.1、创建数据集的配置文件dataset.yaml

训练COCO数据集或者是VOC数据集，可以直接使用已经配置好的coco.yaml 和 voc.yaml文件。如果训练别的数据集，则需要模仿coco.yaml 文件写一个自己的wheat.yaml文件保存在data文件夹下。

# train and val datasets (image directory or *.txt file with image paths)  
train: ./datasets/wheat/images/train/  
val: ./datasets/wheat/images/val/  

# number of classes  
nc: 1  

# class names  
names: ['wheat']

2.2、创建标注文件

如果是自己制作的数据集可以使用LabelImg工具,需要注意的是选择YOLO格式生成标注文件。

如果是已经有的数据集，像是kaggle的wheat数据集，则需要进行转换数据格式成yolo的格式。

我们先来看一下yolo的lable文件里面的格式 00b5fefed.txt

数据排布是：<object-class> <x> <y> <width> <height>

其中，object-class是类的索引，后面的4个值都是相对于整张图片的比例。x是ROI中心的x坐标，y是ROI中心的y坐标，width是ROI的宽，height是ROI的高。

Wheat数据集的lable是都在一个train.csv文件里面，并且bbox是绝对位置，这一点和YOLO不一样。

因此，需要进行Box的转换：

x, y, w, h = bbox[0], bbox[1], bbox[2], bbox[3]  
x_center = (x+ w / 2)/width  
y_center = (y+ h / 2) /hight  
w = w/ width  
h = h / hight

转换好数据格式后需要保存到相应的文件夹格式

至此数据准备阶段已经完成。

3 训练和检测

（1）训练coco数据集，命令如下：

python train.py --data  coco.yaml --cfg yolov5s.yaml --weights '' --batch-size 64
                                        yolov5m.yaml                           40
                                        yolov5l.yaml                           24
                                        yolov5x.yaml                           16

（2）训练自己的数据集：如果想加快训练，可以使用官方在coco数据集上训练好的网络接着进行训练官方的预训练模型链接：https://github.com/ultralytics/yolov5/releases/tag/v3.0

将下载好的模型yolov5*.pt保存在weights文件夹下。

python train.py –data wheat.yaml --cfg  yolov5s.yaml --weights yolov5s.pt --batch-size 16

训练完成是yolov5会保存最好的模型best.pt和最后一个模型last.pt在weights文件夹里。

（3）检测图片

用训练好的模型进行检测：

python detect.py --source inference/images/ --weights  best.pt

4 YOLOv5-MobileNetv2

本文以Yolov5的backbone换成mobilenetv2为例，讲解如何替换其他网络结构。首先，需要在models/common.py里，实现MobileNetv2的 bottleneck 和 Pwconv。

1、Mobilenetv2的bottleneck: InvertedResidual

#mobilenet  Bottleneck  InvertedResidual  
class BottleneckMOB(nn.Module):  
    #c1:inp  c2:oup s:stride  expand_ratio:t  
    def __init__(self, c1, c2, s, expand_ratio):  
        super(BottleneckMOB, self).__init__()  
        self.s = s  
        hidden_dim = round(c1 * expand_ratio)  
        self.use_res_connect = self.s == 1 and c1 == c2  
        if expand_ratio == 1:  
            self.conv = nn.Sequential(  
                # dw  
                nn.Conv2d(hidden_dim, hidden_dim, 3, s, 1, groups=hidden_dim, bias=False),  
                nn.BatchNorm2d(hidden_dim),  
                nn.ReLU6(inplace=True),  
                # pw-linear  
                nn.Conv2d(hidden_dim, c2, 1, 1, 0, bias=False),  
                nn.BatchNorm2d(c2),  
            )  
        else:  
            self.conv = nn.Sequential(  
                # pw  
                nn.Conv2d(c1, hidden_dim, 1, 1, 0, bias=False),  
                nn.BatchNorm2d(hidden_dim),  
                nn.ReLU6(inplace=True),  
                # dw  
                nn.Conv2d(hidden_dim, hidden_dim, 3, s, 1, groups=hidden_dim, bias=False),  
                nn.BatchNorm2d(hidden_dim),  
                nn.ReLU6(inplace=True),  
                # pw-linear  
                nn.Conv2d(hidden_dim, c2, 1, 1, 0, bias=False),  
                nn.BatchNorm2d(c2),  
            )  

    def forward(self, x):  
        if self.use_res_connect:  
            return x + self.conv(x)  
        else:  
            return self.conv(x)

2、Pointwise Convolution

class PW_Conv(nn.Module):  
    def __init__(self, c1, c2):  # ch_in, ch_out  
        super(PW_Conv, self).__init__()  
        self.conv = nn.Conv2d(c1, c2, 1, 1, 0, bias=False)  
        self.bn = nn.BatchNorm2d(c2)  
        self.act = nn.ReLU6(inplace=True)  

    def forward(self, x):  
        return self.act(self.bn(self.conv(x)))

接着需要在yolov5的读取模型配置文件的代码（models/yolo.py的parse_model函数）进行修改，使得能够调用到上面的模块，只需修改下面这部分代码。

n = max(round(n * gd), 1) if n > 1 else n  # depth gain  
if m in [nn.Conv2d, Conv, Bottleneck, SPP, DWConv, MixConv2d, Focus, CrossConv, BottleneckCSP, C3, PW_Conv, BottleneckMOB]:  
    c1, c2 = ch[f], args[0]

并且需要在import引用处加入PW_Conv,BottleneckMOB这两个模块。

from models.common import Conv, Bottleneck,SPP, DWConv, Focus, BottleneckCSP, Concat, NMS, autoShape, PW_Conv,BottleneckMOB

然后就是搭建我们的模型配置文件，我在yolov5s.yaml的基础上进行修改，将yolov5s的backbone替换成mobilenetv2，重新建立了一个模型配置文件yolov5-mobilenet.yaml。

# parameters  
nc: 1  # number of classes  
depth_multiple: 0.33  # model depth multiple  
width_multiple: 0.50  # layer channel multiple  

# anchors  
anchors:  
  - [116,90, 156,198, 373,326]  # P5/32  
  - [30,61, 62,45, 59,119]  # P4/16  
  - [10,13, 16,30, 33,23]  # P3/8  

# YOLOv5 backbone: mobilenet v2  
backbone:  
  # [from, number, module, args]  
  [[-1, 1, nn.Conv2d, [32, 3, 2]],  # 0-P1/2   oup, k, s     640  
   [-1, 1, BottleneckMOB, [16, 1, 1]],  # 1-P2/4   oup, s, t 320  
   [-1, 2, BottleneckMOB, [24, 2, 6]],  #                    320  
   [-1, 1, PW_Conv, [256]],  #4  output p3                   160  
   [-1, 3, BottleneckMOB, [32, 2, 6]],  # 3-P3/8             160  
   [-1, 4, BottleneckMOB, [64, 1, 6]],  # 5                  80  
   [-1, 1, PW_Conv, [512]],  #7 output p4  6                 40  
   [-1, 3, BottleneckMOB, [96, 2, 6]],  # 7                  80  
   [-1, 3, BottleneckMOB, [160, 1, 6,]], #                   40  
   [-1, 1, BottleneckMOB, [320, 1, 6,]], #                   40  
   [-1, 1, nn.Conv2d, [1280, 1, 1]],     #                   40  
   [-1, 1, SPP, [1024, [5, 9, 13]]],  #11     #              40  
  ]  

# YOLOv5 head  
head:  
  [[-1, 3, BottleneckCSP, [1024, False]],  # 12             40  

   [-1, 1, Conv, [512, 1, 1]],                      #       40  
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],      #       40  
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4-7  #       80  
   [-1, 3, BottleneckCSP, [512, False]],  # 16      #       80  

   [-1, 1, Conv, [256, 1, 1]],                      #       80  
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],      #       160  
   [[-1, 3], 1, Concat, [1]],  # cat backbone P3-4          160  
   [-1, 3, BottleneckCSP, [256, False]],            #       160  
   [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1]],  # 21 (P3/8-small)   #        160  

   [-2, 1, Conv, [256, 3, 2]],                     #       160  
   [[-1, 17], 1, Concat, [1]],  # cat head P4      #       160  
   [-1, 3, BottleneckCSP, [512, False]],           #       160  
   [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1]],  # 25 (P4/16-medium)  #       160  

   [-2, 1, Conv, [512, 3, 2]],                     #       160  
   [[-1, 13], 1, Concat, [1]],  # cat head P5-13   #      160  
   [-1, 3, BottleneckCSP, [1024, False]],          #      160  
   [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1]],  # 29 (P5/32-large)           160  

   [[21, 25, 29], 1, Detect, [nc, anchors]],  # Detect(P5, P4, P3)     nc:number class, na:number of anchors  
  ]

到这我们就实现了将yolov5的backbone替换成了mobilenetv2。在使用时只需要将网络结构配置参数—cfg修改成 –cfg yolov5-mobilenet.yaml。

训练指令：

python train.py --data coco.yaml --cfg yolov5-mobilenet.yaml--weights '' --batch-size 64

5 Next
下一篇，《YOLOv5从入门到部署》系列将会介绍：YOLOv5的部署。

发布于 2020-12-14

转自：https://www.zhihu.com/search?type=content&q=%E6%80%8E%E4%B9%88%E6%9B%B4%E6%94%B9yolov5

Ring__Rain

关注

4
点赞
踩
33

收藏

觉得还不错? 一键收藏
0
评论
复制链接

分享到 QQ

分享到新浪微博

扫一扫

专栏目录