58 lines
1.6 KiB
Python
58 lines
1.6 KiB
Python
from ultralytics import YOLO
|
|
import torch
|
|
# 加载预训练模型
|
|
model = YOLO('models/yolov8n-pose.pt')
|
|
# model = YOLO('models/yolo11x-pose.pt')
|
|
# model = YOLO('models/yolov8n-pose.pt')
|
|
# model = YOLO('runs/pose/train2/weights/last.pt')
|
|
|
|
|
|
# 训练模型
|
|
# data:数据集路径
|
|
# epochs:训练轮数
|
|
# imgsz: 图片大小
|
|
# batch: 批次大小
|
|
# device: 使用设备 0:GPU 'cpu':CPU
|
|
# workers: 使用的进程数
|
|
# verbose: 训练进度条
|
|
# resume: 恢复训练
|
|
# name: 训练结果保存名称
|
|
# plot: 绘制训练曲线
|
|
# save: 保存训练结果
|
|
# save_period: 每隔多少轮保存一次模型
|
|
# save_dir: 训练结果保存路径
|
|
# weights: 预训练模型路径
|
|
model.info()
|
|
model.train(data='./dataset1/train.yaml', epochs=300, imgsz=640, batch=32, device=0)
|
|
|
|
|
|
|
|
#
|
|
# # 检查MPS可用性
|
|
# print(f"MPS available: {torch.backends.mps.is_available()}")
|
|
# print(f"MPS built: {torch.backends.mps.is_built()}")
|
|
#
|
|
# model = YOLO('models/yolov8n-pose.pt')
|
|
#
|
|
# # MPS优化配置
|
|
# model.train( data='./dataset1/train.yaml', epochs=300,
|
|
# imgsz=320, # M2 Pro上建议减小尺寸
|
|
# batch=8, # 根据内存调整
|
|
# device='mps', # 使用Apple Metal Performance Shaders
|
|
# workers=2, # MPS下建议2个worker
|
|
# patience=50,
|
|
# lr0=0.01,
|
|
# lrf=0.01,
|
|
# momentum=0.9, # MPS上动量稍小
|
|
# weight_decay=0.0005,
|
|
# warmup_epochs=5.0,
|
|
# box=7.5,
|
|
# pose=1.0, # 增加姿态损失权重
|
|
# kobj=1.5,
|
|
# save=True,
|
|
# exist_ok=True,
|
|
# verbose=True,
|
|
# amp=False # MPS上关闭自动混合精度
|
|
# )
|
|
# print("训练完成")
|