训练之前需要自己先标注,这里推荐在线标注工具,方便快捷。
Make Sense - 在线图像标注 - BimAnt
创建一个文件夹,包含如下子目录,images存放训练(train)和验证图片(val),一般是8:2的数量比例,labels存放对应的yolo格式内容的标注文件(与图片同名.txt)。
创建一个.yaml格式的配置文件,如Goods.yaml
path: D:\Desktop\Python文件\仿真单件分离系统\Goods # 数据集根路径 train: images/train # 训练集路径 val: images/val # 验证集路径 nc: 1 # 类别数量 names: ['Goods'] # 类别名称训练配置,创建train.py进行模型训练
cpu
from ultralytics import YOLO model = YOLO("yolo8n.pt") # 训练配置 results = model.train( data="Goods.yaml", epochs=100, # 训练轮次 imgsz=448, # 设置输入图片尺寸 batch=20, # 每次训练迭代时使用的图像数量 )gpu
from ultralytics import YOLO if __name__ == '__main__': model = YOLO("yolo8n.pt") # 训练配置 results = model.train( data="Goods.yaml", epochs=200, imgsz=448, batch=100, workers=4, # Windows推荐 1~4,不要8 cache="ram" # 把数据集缓存进内存,大幅减少硬盘反复读取 )运行train.py可以在控制台查看训练进度
验证模型,创建val.py进行模型验证
cpu
from ultralytics import YOLO model = YOLO("runs/detect/train/weights/best.pt") metrics = model.val() # 在验证集评估 # 可视化预测结果 results = model.predict( source="Goods/images/val", save=True, conf=0.5, # 置信度阈值 iou=0.45 # IoU阈值 )gpu
from ultralytics import YOLO if __name__ == '__main__': # 加载训练好的最优权重 model = YOLO("runs/detect/train/weights/best.pt") # 验证集评估 + 关闭多进程 workers=0 避免Windows多进程报错 metrics = model.val(workers=0) # 预测推理并保存结果 results = model.predict( source="Goods/images/val", save=True, conf=0.8, iou=0.45, workers=0 )gpu训练需要完成下列操作:
1.清理旧数据
pip uninstall torch torchvision torchaudio -y pip cache purge
2.安装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(f"PyTorch 版本: {torch.__version__}") print(f"CUDA 是否可用: {torch.cuda.is_available()}") print(f"CUDA 版本: {torch.version.cuda}") print(f"GPU 设备名称: {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(format="engine", device=0) # 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, conf=0.5, verbose=False) # 计算推理耗时(毫秒) 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, f"Inference: {inference_ms:.1f}ms", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(annotated_frame, f"Objects: {detections}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(annotated_frame, f"FPS: {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), interpolation=cv2.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, conf=0.5, verbose=False) # 计算推理耗时(毫秒) 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, f"Inference: {inference_ms:.1f}ms", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(annotated_frame, f"Objects: {detections}", (10, 60), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 2) cv2.putText(annotated_frame, f"FPS: {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), interpolation=cv2.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进行部分运算
附加:同种配置下不同架构的推理速度对比