大家可以先在ultralytics/ultralytics文件夹下新建一个 mytrain.py,然后直接复制下面的代码,所有训练有关的超参数都可以在这个文件中调节,不懂超参数可以忽略,这里我说一下比较重要的几个参数:
import sys # 强制扫描导入使用本地ultralytics这个包
sys.path.append("E:/ultralytics") # 改为文件所在的目录地址
from ultralytics import YOLO
if __name__ == '__main__':
# 加载模型
# model = YOLO("yolov8n.yaml") # 从头开始构建新模型
model = YOLO("yolov8n.pt").Load("yolov8n.pt") # 加载预训练模型(推荐用于训练)
# 训练参数 ---------------------------------------------------------------------------
model.train(
data=r'coco128.yaml',
epochs=20, # (int) 训练的周期数
patience=50, # (int) 等待无明显改善以进行早期停止的周期数
batch=8, # (int) 每批次的图像数量(-1 为自动批处理)
imgsz=640, # (int) 输入图像的大小,整数或w,h
save=True, # (bool) 保存训练检查点和预测结果
save_period=-1, # (int) 每x周期保存检查点(如果小于1则禁用)
cache=False, # (bool) True/ram、磁盘或False。使用缓存加载数据
device='', # (int | str | list, optional) 运行的设备,例如 cuda device=0 或 device=0,1,2,3 或 device=cpu
workers=8, # (int) 数据加载的工作线程数(每个DDP进程)
project='runs/train', # (str, optional) 项目名称
name='exp', # (str, optional) 实验名称,结果保存在'project/name'目录下
exist_ok=False, # (bool) 是否覆盖现有实验
pretrained=True, # (bool | str) 是否使用预训练模型(bool),或从中加载权重的模型(str)
optimizer='SGD', # (str) 要使用的优化器,选择=[SGD,Adam,Adamax,AdamW,NAdam,RAdam,RMSProp,auto]
verbose=True, # (bool) 是否打印详细输出
seed=0, # (int) 用于可重复性的随机种子
deterministic=True, # (bool) 是否启用确定性模式
single_cls=False, # (bool) 将多类数据训练为单类
rect=False, # (bool) 如果mode='train',则进行矩形训练,如果mode='val',则进行矩形验证
cos_lr=False, # (bool) 使用余弦学习率调度器
close_mosaic=0, # (int) 在最后几个周期禁用马赛克增强
resume=False, # (bool) 从上一个检查点恢复训练
amp=True, # (bool) 自动混合精度(AMP)训练,选择=[True, False],True运行AMP检查
fraction=1.0, # (float) 要训练的数据集分数(默认为1.0,训练集中的所有图像)
profile=False, # (bool) 在训练期间为记录器启用ONNX和TensorRT速度
freeze= None, # (int | list, 可选) 在训练期间冻结前 n 层,或冻结层索引列表。
# 分割
overlap_mask=True, # (bool) 训练期间是否应重叠掩码(仅适用于分割训练)
mask_ratio=4, # (int) 掩码降采样比例(仅适用于分割训练)
# 分类
dropout=0.0, # (float) 使用丢弃正则化(仅适用于分类训练)
# 超参数 ----------------------------------------------------------------------------------------------
lr0=0.01, # (float) 初始学习率(例如,SGD=1E-2,Adam=1E-3)
lrf=0.01, # (float) 最终学习率(lr0 * lrf)
momentum=0.937, # (float) SGD动量/Adam beta1
weight_decay=0.0005, # (float) 优化器权重衰减 5e-4
warmup_epochs=3.0, # (float) 预热周期(分数可用)
warmup_momentum=0.8, # (float) 预热初始动量
warmup_bias_lr=0.1, # (float) 预热初始偏置学习率
box=7.5, # (float) 盒损失增益
cls=0.5, # (float) 类别损失增益(与像素比例)
dfl=1.5, # (float) dfl损失增益
pose=12.0, # (float) 姿势损失增益
kobj=1.0, # (float) 关键点对象损失增益
label_smoothing=0.0, # (float) 标签平滑(分数)
nbs=64, # (int) 名义批量大小
hsv_h=0.015, # (float) 图像HSV-Hue增强(分数)
hsv_s=0.7, # (float) 图像HSV-Saturation增强(分数)
hsv_v=0.4, # (float) 图像HSV-Value增强(分数)
degrees=0.0, # (float) 图像旋转(+/- deg)
translate=0.1, # (float) 图像平移(+/- 分数)
scale=0.5, # (float) 图像缩放(+/- 增益)
shear=0.0, # (float) 图像剪切(+/- deg)
perspective=0.0, # (float) 图像透视(+/- 分数),范围为0-0.001
flipud=0.0, # (float) 图像上下翻转(概率)
fliplr=0.5, # (float) 图像左右翻转(概率)
mosaic=1.0, # (float) 图像马赛克(概率)
mixup=0.0, # (float) 图像混合(概率)
copy_paste=0.0, # (float) 分割复制-粘贴(概率)
)
metrics = model.val() # 评估模型在验证集上的性能
#results = model("ultralytics/data/images/bus.jpg") # 预测图像
path = model.export(format="onnx") # 将模型导出为 ONNX 格式
# 来自 @小胡学长 进口 YOLO
yolo():这个参数里面可以写你模型yaml文件的路径,也可以直接写.pt 文件
yolo().load():如果你不加.load(),那就是不使用权重从头训练,加了就是使用权重迁移学习
data:数据集yaml文件的路径epochs训练轮数
batch: batch size大小
cache:设置成True可以加速训练
patience:早停轮数,不想设置早停的话,可以把这个值设置成一个非常大的数
optimizer:优化器,最新版应该是提供了7种,注释有写包含哪几种,还可以设置成auto
close_mosaic:最后多少轮关闭马赛克数据增强,这是一种训练技巧,很有效,一般设置10-30
resume:断点续训,一个最实用的参数,直接在这里写你断了的那轮的last.pt的路径就接上了
single_cls:数据集是单个类别的时候记得开启,不开其实也没事
cos_Ir:是否开启余弦学习率
label_smoothing:数据集标注质量很差的时候可以考虑给这个参数设置个0.1-0.5
device:多GPU训练时这里直接写一个列表,比如[e,1]
精简版训练代码
import sys # 强制扫描导入使用本地ultralytics这个包
sys.path.append("/home/aistudio/ultralytics-v8.1.0") # 改为文件所在的目录地址
from ultralytics import YOLO
if __name__ == '__main__':
# 加载模型
model = YOLO("yolov8n.yaml") # 从头开始构建新模型
# model = YOLO("yolov8n.pt") # 加载预训练模型(推荐用于训练)
# 训练参数 ---------------------------------------------------------------------------
model.train(
data='/home/aistudio/data/data.yaml',
epochs=300, # (int) 训练的周期数
patience=50, # (int) 等待无明显改善以进行早期停止的周期数
batch=16, # (int) 每批次的图像数量(-1 为自动批处理)
imgsz=640, # (int) 输入图像的大小,整数或w,h
save=True, # (bool) 保存训练检查点和预测结果
save_period=-1, # (int) 每x周期保存检查点(如果小于1则禁用)
cache=False, # (bool) True/ram、磁盘或False。使用缓存加载数据
device='0', # (int | str | list, optional) 运行的设备,例如 cuda device=0 或 device=0,1,2,3 或 device=cpu
workers=2, # (int) 数据加载的工作线程数(每个DDP进程)
project='runs/train', # (str, optional) 项目名称
name='exp', # (str, optional) 实验名称,结果保存在'project/name'目录下
exist_ok=False, # (bool) 是否覆盖现有实验
)
path = model.export(format="onnx") # 将模型导出为 ONNX 格式
metrics = model.val() # 评估模型在验证集上的性能
# results = model("ultralytics/data/images/bus.jpg") # 预测图像
# 来自 @小胡学长 进口 YOLO
将YOLOv8模型转换为ONNX格式是一个常见的步骤,特别是在需要将模型部署到不同的平台或使用不同的推理引擎时。
ONNX(Open Neural Network Exchange)是一个开放格式,用于表示机器学习模型,允许模型在不同的框架之间共享。
虽然前面的脚本已经集成了模型转换为ONNX格式的代码,但是我还是想单独说说,毕竟有时候我只是需要转换,并不需要训练。
pip install torch onnx==1.16.1 onnx-simplifier #如果你没安装onnx的话,请先下载。
#如果装了最新版的onnx 报动态链接库问题,用下面的方法解决
https://blog.csdn.net/qq_38702496/article/details/141052666
如果上面的你已经安装过了,可以直接看下面的代码:
yolo export model=yolov8n.pt format=onnx simplify=True imgsz=640 output=yolov8n.onnx
import torch
import onnx
from ultralytics import YOLO
# 加载YOLOv8模型
model = YOLO('./best.pt')
# 设置输入大小
img_size = 640
# 创建一个示例输入
dummy_input = torch.randn(1, 3, img_size, img_size)
# 定义输出文件名
output_onnx_file = 'yolov8.onnx'
# 转换为ONNX格式
torch.onnx.export(
model,
dummy_input,
output_onnx_file,
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
},
opset_version=11
)
print(f'Model has been converted to ONNX format and saved to {output_onnx_file}')
为了提高模型的性能,你可以使用onnx-simplifier来简化ONNX模型。简化后的模型通常更小、更快。
import onnxsim
# 加载ONNX模型
onnx_model = onnx.load('yolov8.onnx')
# 简化模型
simplified_model, check = onnxsim.simplify(onnx_model)
# 检查简化后的模型是否与原始模型一致
assert check, "Simplified model is not equivalent to the original model"
# 保存简化后的模型
onnx.save(simplified_model, 'yolov8_simplified.onnx')
print('Simplified ONNX model saved.')
简化模型:
命令行工具 通过--simplify参数自动简化模型。
Python脚本 需要手动调用onnx-simplifier进行简化。
输入大小:
命令行工具:通过--imgsz参数指定输入大小。
Python脚本:通过脚本中的dummy_input显式指定输入大小。
输出文件名:
命令行工具:通过--output参数指定输出文件名。
Python脚本:通过脚本中的output_onnx_file变量指定输出文件名。
换后的ONNX模型应该进行验证,确保其输出与原始PyTorch模型一致。你可以使用以下代码片段来验证:
import onnxruntime
import numpy as np
# 加载简化后的ONNX模型
ort_session = onnxruntime.InferenceSession('yolov8_simplified.onnx')
# 创建一个示例输入
dummy_input_np = np.random.rand(1, 3, img_size, img_size).astype(np.float32)
# 推理
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input_np}
ort_outs = ort_session.run(None, ort_inputs)
# 使用PyTorch模型进行推理
torch_out = model(dummy_input)
# 比较ONNX和PyTorch的输出
np.testing.assert_allclose(torch_out.cpu().detach().numpy(), ort_outs[0], rtol=1e-03, atol=1e-05)
print("Exported model has been successfully tested with ONNXRuntime!")