嵌入式AI---训练自己的yolov5目标检测模型

李华_

已于 2024-08-30 08:55:15 修改

阅读量1.5k

点赞数 17

文章标签：目标检测 YOLO 人工智能

于 2024-07-26 11:05:55 首次发布

本文链接：https://blog.csdn.net/weixin_45816353/article/details/140695107

版权

文章目录

前言
一、yolov5源码下载
二、数据集制作
三、模型训练
四、模型推理
总结

文章阅读顺序：
（1）训练自己的yolov5目标检测模型（本文）
（2）yolov5模型转化为华为昇腾om模型
 （3）在华为昇腾推理自己的yolov5目标检测模型

前言

由于实际项目需要在嵌入式环境部署一个目标检测模型，故以车辆检测为例，基于常用的轻量级检测算法yolov5s，在3060显卡训练了一个自己的车辆检测模型。源码版本为yolov5 v6.0，采用UA-DETRAC数据集训练。

一、yolov5源码下载

本项目选用yolov5 v6.0。源码下载链接：https://github.com/ultralytics/yolov5/releases
1.打开上面的链接，找到v6.0版本
在这里插入图片描述 2.拖到下面打开Assets

3.下载这两个即可

4.将解压后的源码用pycharm打开，yolov5s.pt文件放到源码目录

（PyTorch环境安装省略，本人采用python3.7 torch1.11.0，缺什么包再pip即可。不用完全一样，环境安装可以参考链接）

二、数据集制作

原始数据集链接：（不建议用，过于臃肿，需要特殊处理）

链接：https://pan.baidu.com/s/1QJUSodZZWaBPKxqgnGuIvg 
提取码：rucf

处理后的数据集链接：（建议使用，已根据项目需要进行特殊处理）

链接: https://pan.baidu.com/s/12Qu-0aVSS48cm0AgXLNhqw?pwd=xmrx
提取码: xmrx

处理数据集的原因：UA-DETRAC原始数据集的图片过多，相邻的几帧图片几乎一模一样，故进行了抽帧处理，每隔几帧抽一张作为新数据集。此外，原始数据集的label标注格式不方便工程使用，新数据集进行了格式处理。
（若不想训练车辆检测模型，可以采用labelimg自己标注数据集，参考链接）
将下载好的数据集放入工程目录
在这里插入图片描述

三、模型训练

在确保pytorch环境以及需要的包安装完成后，只需对项目进行以下几处修改即可开始训练。
1.在data目录下添加配置文件UA-DETRAC.yaml
在这里插入图片描述
UA-DETRAC.yaml填充内容如下：(修改为自己的存储路径)

train: D:/yolov5/yolov5-6.0/UA-DETRAC_train/images/train/ # 16551 images
val: D:/yolov5/yolov5-6.0/UA-DETRAC_train/images/val/  # 4952 images
nc: 4
names: ['car',  'bus',  'van', 'others']

2.在models目录下添加配置文件yolov5s_car.yaml
在这里插入图片描述
yolov5s_car.yaml填充内容如下：

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

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

# YOLOv5 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, C3, [128]],
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 9, C3, [256]],
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, C3, [512]],
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 3, C3, [1024, False]],  # 9
  ]

# YOLOv5 head
head:
  [[-1, 1, Conv, [512, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4
   [-1, 3, C3, [512, False]],  # 13

   [-1, 1, Conv, [256, 1, 1]],
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3
   [-1, 3, C3, [256, False]],  # 17 (P3/8-small)

   [-1, 1, Conv, [256, 3, 2]],
   [[-1, 14], 1, Concat, [1]],  # cat head P4
   [-1, 3, C3, [512, False]],  # 20 (P4/16-medium)

   [-1, 1, Conv, [512, 3, 2]],
   [[-1, 10], 1, Concat, [1]],  # cat head P5
   [-1, 3, C3, [1024, False]],  # 23 (P5/32-large)

   [[17, 20, 23], 1, Detect, [nc, anchors]],  # Detect(P3, P4, P5)
  ]

3.修改train.py
修改train.py的parse_opt函数里的以下配置参数：
预训练权重、模型配置文件、数据集配置文件

parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path')
parser.add_argument('--cfg', type=str, default='models/yolov5s_car.yaml', help='model.yaml path')
parser.add_argument('--data', type=str, default=ROOT / 'data/UA-DETRAC.yaml', help='dataset.yaml path')

原train.py设置的训练轮数为300，为了测试代码能否跑通，以下将训练轮数修改为5。

parser.add_argument('--epochs', type=int, default=5)

然后运行train.py即可，见到这个打印信息就说明开始训练了（耐心等待训练结束即可）
在这里插入图片描述

运行train.py可能会出现一些bug，以下为解决办法：
问题1：提示Downloading https://ultralytics.com/assets/Arial.ttf to C:\Users\li\AppData\Roaming\Ultralytics\Arial.ttf…然后程序异常终止。
这是因为yolo项目需要用到Arial.ttf这个字体，运行代码尝试下载该字体导致了下载失败。有两种解决方法，一种是提前下载好该字体，一种是直接屏蔽字体检查
解决方法1：参考
解决方法2：将utils目录下的plots.py的字体检查代码代码注释掉（方法2无法根除问题，训练期间可能会打印错误信息，但不影响训练）

class Annotator:
    # if RANK in (-1, 0):
    #     check_font()  # download TTF if necessary

    # YOLOv5 Annotator for train/val mosaics and jpgs and detect/hub inference annotations
    def __init__(self, im, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'):

问题2：wandb是用来做训练过程可视化的，但wandb频繁报错，影响了训练，如何关闭wandb？
修改工程中的utils/loggers/wandb/wandb_utils.py

# try:
#     import wandb
#
#     assert hasattr(wandb, '__version__')  # verify package import not local dir
#     if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2') and RANK in [0, -1]:
#         wandb.login(timeout=30)
# except (ImportError, AssertionError):
#     wandb = None
wandb = None  # 将上面代码注释掉，添加这行

修改工程中的utils/loggers/–init–.py

# try:
#     import wandb
#
#     assert hasattr(wandb, '__version__')  # verify package import not local dir
# except (ImportError, AssertionError):
#     wandb = None
wandb = None  # 将上面代码注释掉，添加这行

问题3：提示页面文件太小，无法完成操作
这是因为计算机内存太小，加载数据集时爆内存了，修改train.py的配置参数即可。worker会影响dataloader加载数据集，worker越小，内存占用越小。

parser.add_argument('--workers', type=int, default=1, help='maximum number of dataloader workers')

问题4：提示RuntimeError: Unable to find a valid cuDNN algorithm to run convolution
这是因为计算机GPU显存太小，爆显存了，修改train.py的配置参数即可。batch会影响模型训练收敛，batch越小，显存占用越小。

parser.add_argument('--batch-size', type=int, default=8, help='total batch size for all GPUs')

四、模型推理

训练结束后，可以在runs/train/exp目录看见训练结果（可能有多个exp目录，这是因为运行了多次train.py，进最新的那个），在runs/train/exp/weight目录下可以看到best.pt和last.pt模型文件，我们选用best.pt作为推理模型。
1.将runs/train/exp/weight/best.pt复制到工程目录
在这里插入图片描述
2.修改detect.py参数
将权重文件修改为best.pt，将测试路径修改为测试数据集路径

parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'best.pt', help='model path(s)')
parser.add_argument('--source', type=str, default=ROOT / 'UA-DETRAC_test/images', help='file/dir/URL/glob, 0 for webcam')