RoRD-Layout-Recognation/train.py

# train.py

import argparse
import logging
import os
from datetime import datetime
from pathlib import Path

import torch
from torch.utils.data import DataLoader, ConcatDataset, WeightedRandomSampler
from torch.utils.tensorboard import SummaryWriter

from data.ic_dataset import ICLayoutTrainingDataset
from losses import compute_detection_loss, compute_description_loss
from models.rord import RoRD
from utils.config_loader import load_config, to_absolute_path
from utils.data_utils import get_transform

# 设置日志记录
def setup_logging(save_dir):
    """设置训练日志记录"""
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    
    log_file = os.path.join(save_dir, f'training_{datetime.now().strftime("%Y%m%d_%H%M%S")}.log')
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(levelname)s - %(message)s',
        handlers=[
            logging.FileHandler(log_file),
            logging.StreamHandler()
        ]
    )
    return logging.getLogger(__name__)

# --- (已修改) 主函数与命令行接口 ---
def main(args):
    cfg = load_config(args.config)
    config_dir = Path(args.config).resolve().parent

    data_dir = args.data_dir or str(to_absolute_path(cfg.paths.layout_dir, config_dir))
    save_dir = args.save_dir or str(to_absolute_path(cfg.paths.save_dir, config_dir))
    epochs = args.epochs if args.epochs is not None else int(cfg.training.num_epochs)
    batch_size = args.batch_size if args.batch_size is not None else int(cfg.training.batch_size)
    lr = args.lr if args.lr is not None else float(cfg.training.learning_rate)
    patch_size = int(cfg.training.patch_size)
    scale_range = tuple(float(x) for x in cfg.training.scale_jitter_range)

    logging_cfg = cfg.get("logging", None)
    use_tensorboard = False
    log_dir = None
    experiment_name = None

    if logging_cfg is not None:
        use_tensorboard = bool(logging_cfg.get("use_tensorboard", False))
        log_dir = logging_cfg.get("log_dir", "runs")
        experiment_name = logging_cfg.get("experiment_name", "default")

    if args.disable_tensorboard:
        use_tensorboard = False
    if args.log_dir is not None:
        log_dir = args.log_dir
    if args.experiment_name is not None:
        experiment_name = args.experiment_name

    writer = None
    if use_tensorboard and log_dir:
        log_root = Path(log_dir).expanduser()
        experiment_folder = experiment_name or "default"
        tb_path = log_root / "train" / experiment_folder
        tb_path.parent.mkdir(parents=True, exist_ok=True)
        writer = SummaryWriter(tb_path.as_posix())

    logger = setup_logging(save_dir)

    logger.info("--- 开始训练 RoRD 模型 ---")
    logger.info(f"训练参数: Epochs={epochs}, Batch Size={batch_size}, LR={lr}")
    logger.info(f"数据目录: {data_dir}")
    logger.info(f"保存目录: {save_dir}")
    if writer:
        logger.info(f"TensorBoard 日志目录: {tb_path}")

    transform = get_transform()

    # 读取增强与合成配置
    augment_cfg = cfg.get("augment", {})
    elastic_cfg = augment_cfg.get("elastic", {}) if augment_cfg else {}
    use_albu = bool(elastic_cfg.get("enabled", False))
    albu_params = {
        "prob": elastic_cfg.get("prob", 0.3),
        "alpha": elastic_cfg.get("alpha", 40),
        "sigma": elastic_cfg.get("sigma", 6),
        "alpha_affine": elastic_cfg.get("alpha_affine", 6),
        "brightness_contrast": bool(augment_cfg.get("photometric", {}).get("brightness_contrast", True)) if augment_cfg else True,
        "gauss_noise": bool(augment_cfg.get("photometric", {}).get("gauss_noise", True)) if augment_cfg else True,
    }

    # 构建真实数据集
    real_dataset = ICLayoutTrainingDataset(
        data_dir,
        patch_size=patch_size,
        transform=transform,
        scale_range=scale_range,
        use_albu=use_albu,
        albu_params=albu_params,
    )

    # 读取合成数据配置（程序化 + 扩散）
    syn_cfg = cfg.get("synthetic", {})
    syn_enabled = bool(syn_cfg.get("enabled", False))
    syn_ratio = float(syn_cfg.get("ratio", 0.0))
    syn_dir = syn_cfg.get("png_dir", None)

    syn_dataset = None
    if syn_enabled and syn_dir:
        syn_dir_path = Path(to_absolute_path(syn_dir, config_dir))
        if syn_dir_path.exists():
            syn_dataset = ICLayoutTrainingDataset(
                syn_dir_path.as_posix(),
                patch_size=patch_size,
                transform=transform,
                scale_range=scale_range,
                use_albu=use_albu,
                albu_params=albu_params,
            )
            if len(syn_dataset) == 0:
                syn_dataset = None
        else:
            logger.warning(f"合成数据目录不存在，忽略: {syn_dir_path}")
            syn_enabled = False

    # 扩散生成数据配置
    diff_cfg = syn_cfg.get("diffusion", {}) if syn_cfg else {}
    diff_enabled = bool(diff_cfg.get("enabled", False))
    diff_ratio = float(diff_cfg.get("ratio", 0.0))
    diff_dir = diff_cfg.get("png_dir", None)
    diff_dataset = None
    if diff_enabled and diff_dir:
        diff_dir_path = Path(to_absolute_path(diff_dir, config_dir))
        if diff_dir_path.exists():
            diff_dataset = ICLayoutTrainingDataset(
                diff_dir_path.as_posix(),
                patch_size=patch_size,
                transform=transform,
                scale_range=scale_range,
                use_albu=use_albu,
                albu_params=albu_params,
            )
            if len(diff_dataset) == 0:
                diff_dataset = None
        else:
            logger.warning(f"扩散数据目录不存在，忽略: {diff_dir_path}")
            diff_enabled = False

    logger.info(
        "真实数据集大小: %d%s%s" % (
            len(real_dataset),
            f", 合成(程序)数据集: {len(syn_dataset)}" if syn_dataset else "",
            f", 合成(扩散)数据集: {len(diff_dataset)}" if diff_dataset else "",
        )
    )

    # 验证集仅使用真实数据，避免评价受合成样本干扰
    train_size = int(0.8 * len(real_dataset))
    val_size = max(len(real_dataset) - train_size, 1)
    real_train_dataset, val_dataset = torch.utils.data.random_split(real_dataset, [train_size, val_size])

    # 训练集：可与合成数据集合并（程序合成 + 扩散）
    datasets = [real_train_dataset]
    weights = []
    names = []
    # 收集各源与期望比例
    n_real = len(real_train_dataset)
    n_real = max(n_real, 1)
    names.append("real")
    # 程序合成
    if syn_dataset is not None and syn_enabled and syn_ratio > 0.0:
        datasets.append(syn_dataset)
        names.append("synthetic")
    # 扩散合成
    if diff_dataset is not None and diff_enabled and diff_ratio > 0.0:
        datasets.append(diff_dataset)
        names.append("diffusion")

    if len(datasets) > 1:
        mixed_train_dataset = ConcatDataset(datasets)
        # 计算各源样本数
        counts = [len(real_train_dataset)]
        if syn_dataset is not None and syn_enabled and syn_ratio > 0.0:
            counts.append(len(syn_dataset))
        if diff_dataset is not None and diff_enabled and diff_ratio > 0.0:
            counts.append(len(diff_dataset))
        # 期望比例：real = 1 - (syn_ratio + diff_ratio)
        target_real = max(0.0, 1.0 - (syn_ratio + diff_ratio))
        target_ratios = [target_real]
        if syn_dataset is not None and syn_enabled and syn_ratio > 0.0:
            target_ratios.append(syn_ratio)
        if diff_dataset is not None and diff_enabled and diff_ratio > 0.0:
            target_ratios.append(diff_ratio)
        # 构建每个样本的权重
        per_source_weights = []
        for count, ratio in zip(counts, target_ratios):
            count = max(count, 1)
            per_source_weights.append(ratio / count)
        # 展开到每个样本
        weights = []
        idx = 0
        for count, w in zip(counts, per_source_weights):
            weights += [w] * count
            idx += count
        sampler = WeightedRandomSampler(weights, num_samples=len(mixed_train_dataset), replacement=True)
        train_dataloader = DataLoader(mixed_train_dataset, batch_size=batch_size, sampler=sampler, num_workers=4)
        logger.info(
            f"启用混采: real={target_real:.2f}, syn={syn_ratio:.2f}, diff={diff_ratio:.2f}; 总样本={len(mixed_train_dataset)}"
        )
        if writer:
            writer.add_text(
                "dataset/mix",
                f"enabled=true, ratios: real={target_real:.2f}, syn={syn_ratio:.2f}, diff={diff_ratio:.2f}; "
                f"counts: real_train={len(real_train_dataset)}, syn={len(syn_dataset) if syn_dataset else 0}, diff={len(diff_dataset) if diff_dataset else 0}"
            )
    else:
        train_dataloader = DataLoader(real_train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
        if writer:
            writer.add_text("dataset/mix", f"enabled=false, real_train={len(real_train_dataset)}")

    logger.info(f"训练集大小: {len(train_dataloader.dataset)}, 验证集大小: {len(val_dataset)}")
    if writer:
        writer.add_text("dataset/info", f"train={len(train_dataloader.dataset)}, val={len(val_dataset)}")

    val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
    
    model = RoRD().cuda()
    logger.info(f"模型参数数量: {sum(p.numel() for p in model.parameters()):,}")
    
    optimizer = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=1e-4)
    
    # 添加学习率调度器
    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
        optimizer, mode='min', factor=0.5, patience=5
    )
    
    # 早停机制
    best_val_loss = float('inf')
    patience_counter = 0
    patience = 10

    for epoch in range(epochs):
        # 训练阶段
        model.train()
        total_train_loss = 0
        total_det_loss = 0
        total_desc_loss = 0
        
        for i, (original, rotated, H) in enumerate(train_dataloader):
            original, rotated, H = original.cuda(), rotated.cuda(), H.cuda()
            
            det_original, desc_original = model(original)
            det_rotated, desc_rotated = model(rotated)
            
            det_loss = compute_detection_loss(det_original, det_rotated, H)
            desc_loss = compute_description_loss(desc_original, desc_rotated, H)
            loss = det_loss + desc_loss
            
            optimizer.zero_grad()
            loss.backward()
            
            # 梯度裁剪，防止梯度爆炸
            torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
            
            optimizer.step()
            total_train_loss += loss.item()
            total_det_loss += det_loss.item()
            total_desc_loss += desc_loss.item()

            if writer:
                num_batches = len(train_dataloader) if len(train_dataloader) > 0 else 1
                global_step = epoch * num_batches + i
                writer.add_scalar("train/loss_total", loss.item(), global_step)
                writer.add_scalar("train/loss_det", det_loss.item(), global_step)
                writer.add_scalar("train/loss_desc", desc_loss.item(), global_step)
                writer.add_scalar("train/lr", optimizer.param_groups[0]['lr'], global_step)
            
            if i % 10 == 0:
                logger.info(f"Epoch {epoch+1}, Batch {i}, Total Loss: {loss.item():.4f}, "
                          f"Det Loss: {det_loss.item():.4f}, Desc Loss: {desc_loss.item():.4f}")
        
        avg_train_loss = total_train_loss / len(train_dataloader)
        avg_det_loss = total_det_loss / len(train_dataloader)
        avg_desc_loss = total_desc_loss / len(train_dataloader)
        if writer:
            writer.add_scalar("epoch/train_loss_total", avg_train_loss, epoch)
            writer.add_scalar("epoch/train_loss_det", avg_det_loss, epoch)
            writer.add_scalar("epoch/train_loss_desc", avg_desc_loss, epoch)
        
        # 验证阶段
        model.eval()
        total_val_loss = 0
        total_val_det_loss = 0
        total_val_desc_loss = 0
        
        with torch.no_grad():
            for original, rotated, H in val_dataloader:
                original, rotated, H = original.cuda(), rotated.cuda(), H.cuda()
                
                det_original, desc_original = model(original)
                det_rotated, desc_rotated = model(rotated)
                
                val_det_loss = compute_detection_loss(det_original, det_rotated, H)
                val_desc_loss = compute_description_loss(desc_original, desc_rotated, H)
                val_loss = val_det_loss + val_desc_loss
                
                total_val_loss += val_loss.item()
                total_val_det_loss += val_det_loss.item()
                total_val_desc_loss += val_desc_loss.item()
        
        avg_val_loss = total_val_loss / len(val_dataloader)
        avg_val_det_loss = total_val_det_loss / len(val_dataloader)
        avg_val_desc_loss = total_val_desc_loss / len(val_dataloader)
        
        # 学习率调度
        scheduler.step(avg_val_loss)
        
        logger.info(f"--- Epoch {epoch+1} 完成 ---")
        logger.info(f"训练 - Total: {avg_train_loss:.4f}, Det: {avg_det_loss:.4f}, Desc: {avg_desc_loss:.4f}")
        logger.info(f"验证 - Total: {avg_val_loss:.4f}, Det: {avg_val_det_loss:.4f}, Desc: {avg_val_desc_loss:.4f}")
        logger.info(f"学习率: {optimizer.param_groups[0]['lr']:.2e}")
        if writer:
            writer.add_scalar("epoch/val_loss_total", avg_val_loss, epoch)
            writer.add_scalar("epoch/val_loss_det", avg_val_det_loss, epoch)
            writer.add_scalar("epoch/val_loss_desc", avg_val_desc_loss, epoch)
            writer.add_scalar("epoch/lr", optimizer.param_groups[0]['lr'], epoch)
        
        # 早停检查
        if avg_val_loss < best_val_loss:
            best_val_loss = avg_val_loss
            patience_counter = 0
            
            # 保存最佳模型
            if not os.path.exists(save_dir):
                os.makedirs(save_dir)
            save_path = os.path.join(save_dir, 'rord_model_best.pth')
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'best_val_loss': best_val_loss,
                'config': {
                    'learning_rate': lr,
                    'batch_size': batch_size,
                    'epochs': epochs,
                    'config_path': str(Path(args.config).resolve()),
                }
            }, save_path)
            logger.info(f"最佳模型已保存至: {save_path}")
            if writer:
                writer.add_scalar("checkpoint/best_val_loss", best_val_loss, epoch)
        else:
            patience_counter += 1
            if patience_counter >= patience:
                logger.info(f"早停触发！{patience} 个epoch没有改善")
                break
    
    # 保存最终模型
    save_path = os.path.join(save_dir, 'rord_model_final.pth')
    torch.save({
        'epoch': epochs,
        'model_state_dict': model.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'final_val_loss': avg_val_loss,
        'config': {
            'learning_rate': lr,
            'batch_size': batch_size,
            'epochs': epochs,
            'config_path': str(Path(args.config).resolve()),
        }
    }, save_path)
    logger.info(f"最终模型已保存至: {save_path}")
    logger.info("训练完成！")

    if writer:
        writer.add_scalar("final/val_loss", avg_val_loss, epochs - 1)
        writer.close()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="训练 RoRD 模型")
    parser.add_argument('--config', type=str, default="configs/base_config.yaml", help="YAML 配置文件路径")
    parser.add_argument('--data_dir', type=str, default=None, help="训练数据目录，若未提供则使用配置文件中的路径")
    parser.add_argument('--save_dir', type=str, default=None, help="模型保存目录，若未提供则使用配置文件中的路径")
    parser.add_argument('--epochs', type=int, default=None, help="训练轮数，若未提供则使用配置文件中的值")
    parser.add_argument('--batch_size', type=int, default=None, help="批次大小，若未提供则使用配置文件中的值")
    parser.add_argument('--lr', type=float, default=None, help="学习率，若未提供则使用配置文件中的值")
    parser.add_argument('--log_dir', type=str, default=None, help="TensorBoard 日志根目录，覆盖配置文件中的设置")
    parser.add_argument('--experiment_name', type=str, default=None, help="TensorBoard 实验名称，覆盖配置文件中的设置")
    parser.add_argument('--disable_tensorboard', action='store_true', help="禁用 TensorBoard 日志记录")
    main(parser.parse_args())