VisDrone2019上训练YOLOv5（用ultralytics）

原创已于 2023-07-25 22:49:10 修改

· 1.1k 阅读

1 ·

版权

文章标签：

#YOLO #python #机器学习

于 2023-07-25 21:55:27 首次发布

使用的package library： ultralytics

环境：python3.8, torch=1.7.0

把ultralytics代码git clone 到本地

git clone https://github.com/ultralytics/ultralytics/

本地新建一个train.py,内容写下面这些。

from ultralytics import YOLO

# Load a model
model = YOLO("yolov5n.yaml")  # build a new model from scratch

# Use the model
# model.train(data="VisDrone.yaml", epochs=1,batch=1)  # train the model
model.train(data="coco128.yaml", epochs=1,batch=1)

metrics = model.val()  # evaluate model performance on the validation set

在本地做测试用coco128这128张图片。跑通了，上云端服务器的饿时候在换成VisDrone.yaml这个我们要训练的数据集

40G显存，最佳batch size是25张，显示的实时GPU占用仅仅是11.4G，但是显存占用最高的时候是95%.你需要预留一些空间reserve，防止内存爆掉。