机器视觉开发-使用yolo训练和验证自己的模型

📅 2026/7/6 20:24:06
机器视觉开发-使用yolo训练和验证自己的模型
训练之前需要自己先标注这里推荐在线标注工具方便快捷。Make Sense - 在线图像标注 - BimAnt创建一个文件夹包含如下子目录images存放训练train和验证图片val一般是8:2的数量比例labels存放对应的yolo格式内容的标注文件与图片同名.txt。创建一个.yaml格式的配置文件如Goods.yamlpath: D:\Desktop\Python文件\仿真单件分离系统\Goods # 数据集根路径 train: images/train # 训练集路径 val: images/val # 验证集路径 nc: 1 # 类别数量 names: [Goods] # 类别名称训练配置创建train.py进行模型训练cpufrom ultralytics import YOLO model YOLO(yolo8n.pt) # 训练配置 results model.train( dataGoods.yaml, epochs100, # 训练轮次 imgsz448, # 设置输入图片尺寸 batch20, # 每次训练迭代时使用的图像数量 )gpufrom ultralytics import YOLO if __name__ __main__: model YOLO(yolo8n.pt) # 训练配置 results model.train( dataGoods.yaml, epochs200, imgsz448, batch100, workers4, # Windows推荐 1~4不要8 cacheram # 把数据集缓存进内存大幅减少硬盘反复读取 )运行train.py可以在控制台查看训练进度验证模型创建val.py进行模型验证cpufrom ultralytics import YOLO model YOLO(runs/detect/train/weights/best.pt) metrics model.val() # 在验证集评估 # 可视化预测结果 results model.predict( sourceGoods/images/val, saveTrue, conf0.5, # 置信度阈值 iou0.45 # IoU阈值 )gpufrom ultralytics import YOLO if __name__ __main__: # 加载训练好的最优权重 model YOLO(runs/detect/train/weights/best.pt) # 验证集评估 关闭多进程 workers0 避免Windows多进程报错 metrics model.val(workers0) # 预测推理并保存结果 results model.predict( sourceGoods/images/val, saveTrue, conf0.8, iou0.45, workers0 )gpu训练需要完成下列操作1.清理旧数据pip uninstall torch torchvision torchaudio -y pip cache purge2.安装torch torchvision torchaudio以5070ti为例pip install torch torchvision torchaudio -f https://mirrors.aliyun.com/pytorch-wheels/cu121/ -i https://pypi.tuna.tsinghua.edu.cn/simple --trusted-host mirrors.aliyun.com --trusted-host pypi.tuna.tsinghua.edu.cn --timeout 1200运行test.py代码import torch print(fPyTorch 版本: {torch.__version__}) print(fCUDA 是否可用: {torch.cuda.is_available()}) print(fCUDA 版本: {torch.version.cuda}) print(fGPU 设备名称: {torch.cuda.get_device_name(0)})显示如下即可最后在网上随便找张包裹图片进行验证如果需要gpu推理需要安装核心库以TensorRT为例子​ pip install ultralytics -i https://pypi.tuna.tsinghua.edu.cn/simple pip install tensorrt-cu12 -i https://pypi.tuna.tsinghua.edu.cn/simple pip install onnx onnxruntime-gpu opencv-python -i https://pypi.tuna.tsinghua.edu.cn/simple ​gpu推理代码import cv2 import time from ultralytics import YOLO # 1. best.pt 导出为 TensorRT 引擎 print(正在导出 TensorRT 引擎首次运行需要时间...) model YOLO(best.pt) model.export(formatengine, device0) # 2. 加载 TensorRT 引擎进行推理 print(加载 TensorRT 引擎...) trt_model YOLO(best.engine) # 3. 打开视频 cap cv2.VideoCapture(test.avi) if not cap.isOpened(): print(无法打开视频文件) exit() print(开始实时推理TensorRT模式按 Q 退出) # 窗口缩放比例 SCALE 0.5 # 缩小到一半 # FPS 统计 fps_counter 0 fps_display 0 last_fps_time time.time() frame_count 0 while True: ret, frame cap.read() if not ret: break frame_count 1 # 记录推理开始时间 start_time time.perf_counter() # TensorRT 推理 results trt_model(frame, conf0.5, verboseFalse) # 计算推理耗时毫秒 inference_ms (time.perf_counter() - start_time) * 1000 # 绘制结果 annotated_frame results[0].plot() # 计算 FPS fps_counter 1 if time.time() - last_fps_time 1.0: fps_display fps_counter fps_counter 0 last_fps_time time.time() # 显示性能信息 detections len(results[0].boxes) # 显示推理耗时和FPS cv2.putText(annotated_frame, fInference: {inference_ms:.1f}ms, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(annotated_frame, fObjects: {detections}, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(annotated_frame, fFPS: {fps_display}, (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) # 显示 TensorRT 模式标识 cv2.putText(annotated_frame, TensorRT MODE, (annotated_frame.shape[1] - 180, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2) # 画面缩放 height, width annotated_frame.shape[:2] new_width int(width * SCALE) new_height int(height * SCALE) scaled_frame cv2.resize(annotated_frame, (new_width, new_height), interpolationcv2.INTER_AREA) # 显示 cv2.imshow(YOLOv8 TensorRT, scaled_frame) if cv2.waitKey(1) 0xFF ord(q): break cap.release() cv2.destroyAllWindows() print(f处理完成共 {frame_count} 帧)cpu推理代码import cv2 import time from ultralytics import YOLO # 1. 加载 PyTorch 模型 (CPU 推理) print(加载 PyTorch 模型 (CPU 推理)...) model YOLO(best.pt) # 强制使用 CPU不写这句默认gpu如果安装了cuda model.to(cpu) # 2. 打开视频 cap cv2.VideoCapture(test.avi) if not cap.isOpened(): print(无法打开视频文件) exit() print(开始实时推理CPU模式按 Q 退出) # 窗口缩放比例 SCALE 0.5 # 缩小到一半 # FPS 统计 fps_counter 0 fps_display 0 last_fps_time time.time() frame_count 0 while True: ret, frame cap.read() if not ret: break frame_count 1 # 记录推理开始时间 start_time time.perf_counter() # CPU 推理 results model(frame, conf0.5, verboseFalse) # 计算推理耗时毫秒 inference_ms (time.perf_counter() - start_time) * 1000 # 绘制结果 annotated_frame results[0].plot() # 计算 FPS fps_counter 1 if time.time() - last_fps_time 1.0: fps_display fps_counter fps_counter 0 last_fps_time time.time() # 显示性能信息 detections len(results[0].boxes) # 显示推理耗时和FPS cv2.putText(annotated_frame, fInference: {inference_ms:.1f}ms, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(annotated_frame, fObjects: {detections}, (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(annotated_frame, fFPS: {fps_display}, (10, 90), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) # 显示 CPU 模式标识 cv2.putText(annotated_frame, CPU MODE, (annotated_frame.shape[1] - 150, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 0, 255), 2) # 画面缩放 height, width annotated_frame.shape[:2] new_width int(width * SCALE) new_height int(height * SCALE) scaled_frame cv2.resize(annotated_frame, (new_width, new_height), interpolationcv2.INTER_AREA) # 显示 cv2.imshow(YOLO PyTorch, scaled_frame) if cv2.waitKey(1) 0xFF ord(q): break cap.release() cv2.destroyAllWindows() print(f处理完成共 {frame_count} 帧)推理速度对比注意安装cuda后PyTorch会检测CUDA环境即使选择CPU推理也会默认启用GPU进行部分运算附加同种配置下不同架构的推理速度对比