建议使用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
创建 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 )
需要注意,预测图片的路径不能包含中文,否则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 # 是否优化模型
)
将模型导出为onnx
yolo task=detect mode=export model=.\runs\detect\train\weights\best.pt format=onnx opset=12
然后使用python quantize_dynamic
进行量化处理,推荐直接使用opt_onnx_model.py一步处理