train_attention_yolo.py

import os
from models.yolov8_depthwise_attention import YOLOv8DepthwiseAttention

# 配置路径
DATASET_PATH = './DAWN_YOLO_AUGMENTED'
LABEL_DIR = os.path.join(DATASET_PATH, 'labels/train')
YAML_PATH = os.path.join(DATASET_PATH, 'data.yaml')
MAX_CLASS = 8

def train_attention_yolo():
    print("🚀 开始训练带注意力机制的YOLOv8模型...")
    
    # 创建目录结构
    os.makedirs('models', exist_ok=True)
    
    # 创建模型
    model = YOLOv8DepthwiseAttention('weights/yolov8n.pt')
    
    # 计算参数量
    params = sum(p.numel() for p in model.model.parameters())
    print(f"✅ 模型参数量: {params/1e6:.2f}M")
    
    # 开始训练
    model.train(
        data=YAML_PATH,
        epochs=100,  # 100-150轮
        imgsz=640,
        batch=8,  # 批次大小
        lr0=0.001,  # 初始学习率
        optimizer="AdamW",  # 优化器
        weight_decay=0.01,  # 权重衰减
        cos_lr=True,  # 余弦退火学习率
        patience=15,  # 早停机制
        name="dawn_train_attention",
        save=True
    )
    print("✅ 训练完成！结果保存在 runs/detect/dawn_train_attention/")

if __name__ == '__main__':
    train_attention_yolo()