scripts/train.py

#!/usr/bin/env python3
"""
Lumina训练脚本
用于训练轻量级图像生成模型
"""

import os
import sys
import argparse
import yaml
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
import warnings

# 添加项目根目录到Python路径
sys.path.append(os.path.dirname(os.path.dirname(__file__)))

from src.models.unet_light import UNetLight
from src.models.diffusion import DiffusionProcess, DiffusionModel
from src.data.dataset import create_data_loaders
from src.data.text_encoder import create_text_encoder
from src.training.trainer_p4 import P4Trainer
from src.training.memory_manager import MemoryOptimizer
from src.training.callbacks import create_default_callbacks


def load_config(config_path: str) -> dict:
    """加载配置文件"""
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    return config


def setup_environment(config: dict):
    """设置训练环境"""
    # 设置随机种子
    seed = config.get('seed', 42)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
    
    # 设置CUDA设备
    device = config.get('device', 'cuda' if torch.cuda.is_available() else 'cpu')
    if device == 'cuda' and not torch.cuda.is_available():
        warnings.warn("CUDA不可用，使用CPU")
        device = 'cpu'
    
    # 创建输出目录
    output_dir = config.get('output_dir', './output')
    os.makedirs(output_dir, exist_ok=True)
    
    # 设置日志
    log_dir = config.get('log_dir', './logs')
    os.makedirs(log_dir, exist_ok=True)
    
    print(f"环境设置完成:")
    print(f"  设备: {device}")
    print(f"  随机种子: {seed}")
    print(f"  输出目录: {output_dir}")
    print(f"  日志目录: {log_dir}")
    
    return device


def create_model(config: dict, device: torch.device) -> nn.Module:
    """创建模型"""
    # 加载模型配置
    model_config_path = config.get('model_config', 'configs/model/unet_light.yaml')
    model_config = load_config(model_config_path)
    
    # 创建UNet模型
    model = UNetLight(model_config)
    
    # 加载预训练权重（如果有）
    pretrained_path = config.get('pretrained_path')
    if pretrained_path and os.path.exists(pretrained_path):
        print(f"加载预训练权重: {pretrained_path}")
        checkpoint = torch.load(pretrained_path, map_location='cpu')
        model.load_state_dict(checkpoint['model_state_dict'])
    
    # 移动到设备
    model = model.to(device)
    
    # 打印模型信息
    total_params = sum(p.numel() for p in model.parameters())
    trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
    
    print(f"模型创建完成:")
    print(f"  总参数量: {total_params:,}")
    print(f"  可训练参数量: {trainable_params:,}")
    print(f"  模型大小: {total_params * 4 / 1024**2:.2f} MB (fp32)")
    
    return model


def create_diffusion(config: dict) -> DiffusionProcess:
    """创建扩散过程"""
    diffusion_config_path = config.get('diffusion_config', 'configs/model/diffusion.yaml')
    diffusion_config = load_config(diffusion_config_path)
    
    diffusion = DiffusionProcess(diffusion_config)
    
    print(f"扩散过程创建完成:")
    print(f"  训练时间步: {diffusion.num_train_timesteps}")
    print(f"  推理时间步: {diffusion.num_inference_timesteps}")
    print(f"  Beta调度: {diffusion.beta_schedule}")
    
    return diffusion


def create_data_pipeline(config: dict):
    """创建数据管道"""
    data_config_path = config.get('data_config', 'configs/data/laion_filtered.yaml')
    data_config = load_config(data_config_path)
    
    # 创建文本编码器
    text_encoder = create_text_encoder(data_config)
    
    # 创建数据加载器
    train_loader, val_loader = create_data_loaders(data_config)
    
    print(f"数据管道创建完成:")
    print(f"  训练集大小: {len(train_loader.dataset)}")
    print(f"  验证集大小: {len(val_loader.dataset) if val_loader else 0}")
    print(f"  批次大小: {train_loader.batch_size}")
    print(f"  梯度累积步数: {config.get('gradient_accumulation_steps', 8)}")
    
    return train_loader, val_loader, text_encoder


def create_optimizer(model: nn.Module, config: dict):
    """创建优化器"""
    optimizer_config = config.get('optimizer', {})
    optimizer_type = optimizer_config.get('type', 'AdamW')
    learning_rate = optimizer_config.get('learning_rate', 1e-4)
    weight_decay = optimizer_config.get('weight_decay', 0.01)
    
    if optimizer_type == 'AdamW':
        optimizer = torch.optim.AdamW(
            model.parameters(),
            lr=learning_rate,
            weight_decay=weight_decay,
            betas=(0.9, 0.999),
            eps=1e-8
        )
    elif optimizer_type == 'Adam':
        optimizer = torch.optim.Adam(
            model.parameters(),
            lr=learning_rate,
            weight_decay=weight_decay
        )
    else:
        raise ValueError(f"未知的优化器类型: {optimizer_type}")
    
    print(f"优化器创建完成:")
    print(f"  类型: {optimizer_type}")
    print(f"  学习率: {learning_rate}")
    print(f"  权重衰减: {weight_decay}")
    
    return optimizer


def setup_memory_optimization(model: nn.Module, optimizer, config: dict):
    """设置内存优化"""
    memory_optimizer = MemoryOptimizer(config)
    memory_optimizer.setup_model_optimizations(model, optimizer)
    
    # 打印内存信息
    if torch.cuda.is_available():
        allocated = torch.cuda.memory_allocated() / 1024**3
        reserved = torch.cuda.memory_reserved() / 1024**3
        print(f"内存优化设置完成:")
        print(f"  GPU已分配: {allocated:.2f} GB")
        print(f"  GPU已保留: {reserved:.2f} GB")
    
    return memory_optimizer


def train(config_path: str, resume_from: str = None):
    """训练主函数"""
    print("=" * 60)
    print("Lumina 训练开始")
    print("=" * 60)
    
    # 加载配置
    config = load_config(config_path)
    
    # 设置环境
    device = setup_environment(config)
    
    # 创建模型
    model = create_model(config, device)
    
    # 创建扩散过程
    diffusion = create_diffusion(config)
    
    # 创建扩散模型
    diffusion_model = DiffusionModel(model, diffusion)
    
    # 创建数据管道
    train_loader, val_loader, text_encoder = create_data_pipeline(config)
    
    # 创建优化器
    optimizer = create_optimizer(model, config)
    
    # 设置内存优化
    memory_optimizer = setup_memory_optimization(model, optimizer, config)
    
    # 创建训练器
    trainer = P4Trainer(
        model=model,
        diffusion=diffusion,
        optimizer=optimizer,
        train_loader=train_loader,
        val_loader=val_loader,
        config=config,
        device=device
    )
    
    # 创建回调
    callbacks = create_default_callbacks(config)
    
    # 加载检查点（如果存在）
    if resume_from and os.path.exists(resume_from):
        print(f"从检查点恢复训练: {resume_from}")
        trainer.load_checkpoint(resume_from)
    
    # 开始训练
    try:
        print("\n开始训练...")
        trainer.train()
        
        print("\n" + "=" * 60)
        print("训练完成!")
        print(f"最佳验证损失: {trainer.best_loss:.4f}")
        print(f"总训练步数: {trainer.global_step}")
        print("=" * 60)
        
    except KeyboardInterrupt:
        print("\n训练被中断")
    except Exception as e:
        print(f"\n训练出错: {e}")
        import traceback
        traceback.print_exc()
    
    finally:
        # 保存最终检查点
        final_checkpoint = os.path.join(
            config.get('checkpoint_dir', './checkpoints'),
            'final_model.pt'
        )
        trainer.save_checkpoint(final_checkpoint)


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="训练Lumina图像生成模型")
    
    parser.add_argument(
        "--config",
        type=str,
        default="configs/training/p4_optimized.yaml",
        help="训练配置文件路径"
    )
    
    parser.add_argument(
        "--resume",
        type=str,
        help="从检查点恢复训练"
    )
    
    parser.add_argument(
        "--debug",
        action="store_true",
        help="调试模式"
    )
    
    args = parser.parse_args()
    
    # 调试模式设置
    if args.debug:
        import warnings
        warnings.filterwarnings("always")
        torch.autograd.set_detect_anomaly(True)
        print("调试模式已启用")
    
    # 开始训练
    train(args.config, args.resume)


if __name__ == "__main__":
    main()