from ultralytics import YOLO
import torch

# 检查GPU信息
print(f"CUDA可用: {torch.cuda.is_available()}")
print(f"GPU数量: {torch.cuda.device_count()}")
print(f"当前GPU: {torch.cuda.get_device_name(0)}")
print(f"GPU内存: {torch.cuda.get_device_properties(0).total_memory / 1024 ** 3:.1f} GB")

model = YOLO('models/yolov8n-pose.pt')

# 3080优化配置
model.train(
    data='./dataset1/train.yaml',
    epochs=300,
    imgsz=640,  # 3080可以处理原尺寸
    batch=32,  # 根据12GB内存调整
    device=0,  # 使用GPU 0
    workers=8,  # 充分利用CPU核心
    patience=50,
    # lr0=0.01,
    # lrf=0.01,
    # momentum=0.937,
    # weight_decay=0.0005,
    # warmup_epochs=3.0,
    #
    # # 损失权重调整（重点解决姿态问题）
    # box=7.5,
    # pose=1.5,  # 提高姿态损失权重
    # kobj=2.0,  # 提高关键点目标权重
    # cls=1.0,
    #
    # # 性能优化
    # amp=True,  # 自动混合精度训练
    # cos_lr=True,  # 余弦学习率调度
    # close_mosaic=10,  # 最后10epoch关闭马赛克增强
    #
    # # 数据增强
    # hsv_h=0.015,
    # hsv_s=0.7,
    # hsv_v=0.4,
    # fliplr=0.5,  # 水平翻转对姿态很重要

    save=True,
    exist_ok=True,
    verbose=True
)
print("训练完成")