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

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)

    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}")

    transform = get_transform()

    dataset = ICLayoutTrainingDataset(
        data_dir,
        patch_size=patch_size,
        transform=transform,
        scale_range=scale_range,
    )
    
    logger.info(f"数据集大小: {len(dataset)}")
    
    # 分割训练集和验证集
    train_size = int(0.8 * len(dataset))
    val_size = len(dataset) - train_size
    train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, val_size])
    
    logger.info(f"训练集大小: {len(train_dataset)}, 验证集大小: {len(val_dataset)}")
    
    train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    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 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)
        
        # 验证阶段
        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 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}")
        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 __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="学习率，若未提供则使用配置文件中的值")
    main(parser.parse_args())