|
|
| """
|
| Lumina训练脚本
|
| 用于训练轻量级图像生成模型
|
| """
|
|
|
| import os
|
| import sys
|
| import argparse
|
| import yaml
|
| import torch
|
| import torch.nn as nn
|
| from torch.utils.data import DataLoader
|
| import warnings
|
|
|
|
|
| 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)
|
|
|
|
|
| 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)
|
|
|
|
|
| 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() |
