from ultralytics import YOLO

if __name__ == '__main__':
    # ✅ 推荐：使用官方预训练分割模型
    model = YOLO(r'ultralytics/cfg/models/11/yolo11-seg.yaml')

    # 开始训练
    results = model.train(
        data='data.yaml',           # 数据配置文件
        imgsz=1280,                  # 输入图像大小
        epochs=100,                 # 训练轮数
        batch=4,                   # 每批图像数量
        workers=10,                 # 数据加载线程数
        device='0',                 # 使用 GPU 0
        project='runs/train',       # 保存项目目录
        name='exp',                 # 实验名称
        exist_ok=False,             # 不覆盖已有实验
        optimizer='AdamW',          # 可选优化器
        lr0=0.001,                  # 初始学习率
        patience=20,                # 早停轮数
    )