YOLOv8 训练用骨架工程

使用 YOLOv8 训练自己的数据集
改造自：https://github.com/prophet-mu/YOLOv8_Train_Guide
原作者相关使用说明：语雀(prophet-mu)
非常感谢 @prophet-mu 的这个仓库带我入门yolo

环境准备

建议使用Anaconda创建虚拟环境，具体安装可以网上搜索或者上B站学习
环境搭建完毕之后安装pytorch和cuda环境，具体参考上面语雀文档中的说明，如果无GPU可以不安装cuda直接用CPU训练也可以只是慢很多。
conda创建虚拟环境，用于Yolo训练，记得切换

准备训练数据

创建数据集

将数据集保存到data/{识别类型}
- 例如识别蚂蚁森林，识别类型定义为 forest 则保存到 datasets/forest 下面
对于数据集标注，建议使用anylabeling工具进行标注，标注完成后将得到图片和对应的json标注信息

数据集保存完毕后修改 convert_to_yolo_txt.py 然后运行他，将json转换成yolo标注的txt文件

# 初始化训练集时，指定为空数组，将自动统计并打印所有标签
specific_labels = []
# 指定原始目录和保存目录 保存目录为datasets下面
root_path = './data/forest'
target_path = './datasets/forest'

运行完毕后将打印出当前数据集中的所有标签信息，复制保存后面会用到

分割数据集为：train\val\test三个部分

修改 split.py 然后运行它，将 datasets/forest 中的数据集进行拆分

root_path = "./datasets"
source_train_type = 'forest'
target_train_type = 'forest'

创建配置文件

在config目录下创建自己的训练配置文件，例如forest.yaml

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
# 数据集所在位置，建议使用绝对路径
path: K:/YOLOV8_train_clean/datasets/forest
train: # train images (relative to 'path')  16551 images
  - ./train
val: # val images (relative to 'path')  4952 images
  ./val
test: # test images (optional)
  ./test
# Classes
# 设置标签值，注意如果需要持续训练优化模型，这些标签的顺序不能修改，否则需要重头开始训练。新增标签可以一直往后增加
names:
  0: cannot
  1: collect
  2: countdown
  3: help_revive
  4: item
  5: tree
  6: water
  7: waterBall
  8: stroll_btn
  9: sea_ball
  10: sea_garbage

开始训练

有两种方式可以进行训练，一种是通过命令行，另一种是python代码选择自己喜欢的方式即可

命令行方式

打开终端，输入如下命令，具体参数说明见官网

yolo task=detect mode=train model=yolov8n.pt data=config/forest.yaml batch=-1 epochs=20

训练完毕后会将模型保存到 runs/detect/train/weights/best.pt 中

如果增加了数据集，需要在原有模型基础上继续训练，则修改命令如下

# 指定model为上一次训练的结果
yolo task=detect mode=train model=runs/detect/train/weights/best.pt data=config/forest.yaml batch=-1 epochs=100

python方式

创建 train.py 添加如下内容：

from ultralytics import YOLO

if __name__ == '__main__':
    # Load a model
    model = YOLO(r'yolov8n.pt')  # load a pretrained model (recommended for training)

    # Train the model
    results = model.train(data='./config/forest.yaml', epochs=20, imgsz=640, device=0, batch=-1)

训练完毕后会将模型保存到 runs/detect/train/weights/best.pt 中

如果增加了数据集，需要在原有模型基础上继续训练，则修改代码如下

from ultralytics import YOLO

if __name__ == '__main__':
    # 指定上一次训练的结果
    model = YOLO('runs/detect/train/weights/best.pt')

    # Train the model
    results = model.train(data='./config/forest.yaml', epochs=20, imgsz=640, device=0, batch=-1)

使用训练后模型预测

如果直接在本机预测，可以直接使用pt模型

需要注意，预测图片的路径不能包含中文，否则opencv无法正确读取

from ultralytics import YOLO  
best_model = YOLO(model='K:/YOLOV8_train_clean/runs/detect/train/weights/best.pt')
# 指定project為當前執行目錄，否則會按settings文件中的地址進行保存
result = best_model.predict(project=project_path, source=img_path, save=False)
# 后续针对result进行处理即可，具体参考auto_predict.py即可

模型导出并量化

跨平台使用时，可以使用onnxruntime推理onnx模型，支持c++和java等
量化模型可以减小模型体积，加快模型推理速度，但是会损失一定的精度

修改并运行 opt_onnx_model.py 即可将pt模型转换成onnx模型并量化

if __name__ == '__main__':
  # 指定训练后的模型路径
  model_path = 'runs/detect/manor_v4/weights/best.pt'
  onnx_path = model_path.replace('.pt', '.onnx')
  model_quant_dynamic = onnx_path.replace('.onnx', '_lite.onnx')
  # 导出onnx模型
  export_onnx_if_not_exists()
  # 动态量化
  quantize_dynamic(
      model_input=onnx_path,  # 输入模型
      model_output=model_quant_dynamic,  # 输出模型
      weight_type=QuantType.QUInt8,  # 参数类型 Int8 / UInt8
      optimize_model=True  # 是否优化模型
  )

也可以直接使用yolo命令导出

将模型导出为onnx
yolo task=detect mode=export model=.\runs\detect\train\weights\best.pt format=onnx opset=12
然后使用python quantize_dynamic 进行量化处理，推荐直接使用opt_onnx_model.py一步处理

TonyJiangWJ / Yolov8_Train