from ultralytics import YOLO
import torch


# 检查CUDA是否可用
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using device: {device}")

# 加载模型
model = YOLO("runs/detect/train6/weights/last.pt").to(device)

# 设置新的分辨率
imgsz = 1024  # 这里将图像尺寸调整为 1280x1280（你可以根据显存调整尺寸）

# 训练模型，传入增强参数
model.train(
    data="dataset/dataset.yaml",  # 你的数据集配置文件
    epochs=1000,  # 训练轮次
    imgsz=imgsz,  # 使用更高的分辨率
    device=[1],  # 使用第一块 GPU（如果有多个 GPU，可以调整）
    hsv_v=0.3,  # 修改图像亮度的一部分，帮助模型在不同光照条件下表现良好
    cos_lr=True,  # 启用余弦学习率调度
    batch = -1,  # 自动调整批量大小以适应显存
)