/Yolov8_Train

yolov8训练

Primary LanguagePython

YOLOv8 训练用骨架工程

环境准备

  • 建议使用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一步处理