模型训练与验证模版

1.模型训练
from ultralytics import YOLO
import torch

if name == 'main':
# # 打印 PyTorch 版本
# print(f"PyTorch version: {torch.version}")
# # 检查 CUDA 是否可用
# cuda_available = torch.cuda.is_available()
# print(f"CUDA available: {cuda_available}")

exit()

model = YOLO("yolov8s.pt")

model = YOLO("runs/detect/train/weights/last.pt")
model.train(data=r"D:\01-AI\20260914\dataset\images\images\YOLODataset\dataset.yaml",
batch=-1, # 根据显存调整，显存不够就调小
imgsz = 640, # 图片的尺寸 640 1280
epochs = 30, # 训练的轮次 200 300
# lr0 = 0.001, # 降低初始学习率
optimizer = 'SGD', # 优化器 Adam
# patience = 50, # 早停耐心值
# close_mosaic = 20, # 最后 20 轮关闭 Mosaic
# copy_paste = 0.3, # 启用复制粘贴增强
# mixup = 0.0, # 关闭 MixUp
# hsv_s = 0.5, # 饱和度增强适中
# hsv_v = 0.4, # 亮度增强适中
# save_period = 10, # 定期保存
device = 0, # 指定 GPU
workers=1, # 线程数
resume = True
)

2.模型验证
from ultralytics import YOLO
import torch

if name == 'main':

model = YOLO("runs/detect/train/weights/best.pt")
model.val(data=r"D:\01-AI\20260914\dataset\images\images\YOLODataset\dataset.yaml",
save_json=True, # 收集模型的预测置信度和真实标签 -> 绘制 ROC 曲线与 AUC
)