跳转至

训练YoloV5

训练代码

https://github.com/wojiazaiyugang/yolov5

训练环境

docker pull ultralytics/yolov5:latest

ln -sf /opt/conda/bin/python /usr/bin/python

数据集准备

参考 https://python.iitter.com/other/195770.html

voc格式标注。可以用scripts/datasets/generate_yolov5_label.py来将voc格式的数据文件夹生成对应的label文件及目录结构，目录结构为

img.png

然后在配置文件中修改对应的label文件。

path: /mnt/nfs-storage/yujiannan/Workspace/yolov5/data/basket
train: # train images (relative to 'path')  16551 images
  - images/train
val: # val images (relative to 'path')  4952 images
  - images/test
test: # test images (optional)
  - images/test

# Classes
nc: 1  # number of classes
names: ['basket']  # class names

训练后好像数据处会生成.cache缓存文件，如果有的话就删一下

配置文件

见models文件夹，新建个配置，修改一下nc改为要检测的类别数就可以了

训练

python train.py --data data/basket.yaml --cfg models/yolov5s_basket.yaml --weights '' --batch-size -1 --save-period 5

模型转onnx/mlmodel

python export.py --weights runs/train/exp/weights/best.pt --include [onnx, coreml] --simplify --device 0

--device 0不要忘记，否则onnx转tensorrt时会报错

评论