import torch
from ultralytics import YOLO

# Create a new YOLO model from scratch
model = YOLO("yolo12n.yaml")

# 控制 GPU 显存使用（例如限制为总显存的 70%）
#torch.cuda.set_per_process_memory_fraction(0.7, device=0)

# 增加训练参数配置：图像大小、批次大小、学习率、数据增强等
if __name__ == "__main__":
    results = model.train(
        data="data.yaml",
        epochs=50,          # 增加训练轮次到50（线缆识别通常需要更多迭代）
        imgsz=640,          # 设置输入图像尺寸
        batch=4,           # 根据GPU内存调整批次大小
        lr0=0.01,           # 初始学习率
        augment=True,       # 启用数据增强
        name="cable_detection",  # 训练结果保存目录
        device=0,           # 使用GPU 0
        workers=0
)

# 新增：训练完成后自动评估
results = model.val(data="data.yaml")

# 新增：导出最佳模型到PyTorch TorchScript格式（移除ONNX依赖）
success = model.export(format="torchscript")  # 使用支持的PyTorch导出格式

# 新增：保存原始PyTorch模型权重到model.pt
model.save('model.pt')  # 直接保存.pt格式的原始模型文件