train.py

import os
import argparse
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.utils.data import DataLoader, random_split
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm
from dataset import HandwriteDataset
from model import UNet
import torchvision
import torchvision.models as models

class CharbonnierLoss(nn.Module):
    def __init__(self, eps=1e-3):
        super().__init__()
        self.eps = eps

    def forward(self, pred, target):
        diff = pred - target
        loss = torch.sqrt(diff * diff + self.eps * self.eps)
        return loss.mean()

class PerceptualLoss(nn.Module):
    """使用 VGG16 的感知损失"""
    def __init__(self):
        super().__init__()
        vgg = models.vgg16(weights=models.VGG16_Weights.IMAGENET1K_V1).features
        self.slice1 = nn.Sequential(*list(vgg[:4]))   # relu1_2
        self.slice2 = nn.Sequential(*list(vgg[4:9]))  # relu2_2
        self.slice3 = nn.Sequential(*list(vgg[9:16])) # relu3_3
        
        # ImageNet 归一化参数
        self.register_buffer('mean', torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
        self.register_buffer('std', torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))
        
        for param in self.parameters():
            param.requires_grad = False
    
    def normalize(self, x):
        """将 [0, 1] 的图像归一化到 ImageNet 分布"""
        return (x - self.mean) / self.std
    
    def forward(self, pred, target):
        # Convert to the same dtype and device as input to handle mixed precision training
        if pred.device != self.mean.device or pred.device != self.slice1[0].weight.device:
            self.to(pred.device)
        
        # VGG 不支持 half，转回 float32 计算
        pred_float = pred.float()
        target_float = target.float()
        
        # 归一化到 ImageNet 分布
        pred_float = self.normalize(pred_float)
        target_float = self.normalize(target_float)
        
        with torch.amp.autocast('cuda', enabled=False):
            pred_1 = self.slice1(pred_float)
            target_1 = self.slice1(target_float)
            pred_2 = self.slice2(pred_1)
            target_2 = self.slice2(target_1)
            pred_3 = self.slice3(pred_2)
            target_3 = self.slice3(target_2)
            
            loss = F.l1_loss(pred_1, target_1) + \
                   F.l1_loss(pred_2, target_2) + \
                   F.l1_loss(pred_3, target_3)
        return loss

class SSIMLoss(nn.Module):
    """结构相似性损失"""
    def __init__(self, window_size=11):
        super().__init__()
        self.window_size = window_size
        self.channel = 3
        self.window = self.create_window(window_size, self.channel)
    
    def gaussian(self, window_size, sigma):
        gauss = torch.Tensor([
            torch.exp(torch.tensor(-(x - window_size//2)**2 / float(2*sigma**2)))
            for x in range(window_size)
        ])
        return gauss / gauss.sum()
    
    def create_window(self, window_size, channel):
        _1D_window = self.gaussian(window_size, 1.5).unsqueeze(1)
        _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
        window = _2D_window.expand(channel, 1, window_size, window_size).contiguous()
        return window
    
    def forward(self, pred, target):
        if self.window.device != pred.device:
            self.window = self.window.to(pred.device)
        if self.window.dtype != pred.dtype:
            self.window = self.window.to(pred.dtype)
        
        mu1 = F.conv2d(pred, self.window, padding=self.window_size//2, groups=self.channel)
        mu2 = F.conv2d(target, self.window, padding=self.window_size//2, groups=self.channel)
        
        mu1_sq = mu1.pow(2)
        mu2_sq = mu2.pow(2)
        mu1_mu2 = mu1 * mu2
        
        sigma1_sq = F.conv2d(pred * pred, self.window, padding=self.window_size//2, groups=self.channel) - mu1_sq
        sigma2_sq = F.conv2d(target * target, self.window, padding=self.window_size//2, groups=self.channel) - mu2_sq
        sigma12 = F.conv2d(pred * target, self.window, padding=self.window_size//2, groups=self.channel) - mu1_mu2
        
        C1 = 0.01 ** 2
        C2 = 0.03 ** 2
        
        ssim_map = ((2 * mu1_mu2 + C1) * (2 * sigma12 + C2)) / \
                   ((mu1_sq + mu2_sq + C1) * (sigma1_sq + sigma2_sq + C2))
        
        # 数值稳定性：裁剪 SSIM 值到合理范围
        ssim_map = torch.clamp(ssim_map, -1, 1)
        return 1 - ssim_map.mean()

class EdgeLoss(nn.Module):
    """边缘保持损失"""
    def __init__(self):
        super().__init__()
        # Sobel 算子
        self.sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3)
        self.sobel_y = torch.tensor([[-1, -2, -1], [0, 0, 0], [1, 2, 1]], dtype=torch.float32).view(1, 1, 3, 3)
    
    def forward(self, pred, target):
        if self.sobel_x.device != pred.device:
            self.sobel_x = self.sobel_x.to(pred.device)
            self.sobel_y = self.sobel_y.to(pred.device)
        if self.sobel_x.dtype != pred.dtype:
            self.sobel_x = self.sobel_x.to(pred.dtype)
            self.sobel_y = self.sobel_y.to(pred.dtype)
        
        # 转换为灰度图
        pred_gray = 0.299 * pred[:, 0:1] + 0.587 * pred[:, 1:2] + 0.114 * pred[:, 2:3]
        target_gray = 0.299 * target[:, 0:1] + 0.587 * target[:, 1:2] + 0.114 * target[:, 2:3]
        
        # 计算边缘
        pred_edge_x = F.conv2d(pred_gray, self.sobel_x, padding=1)
        pred_edge_y = F.conv2d(pred_gray, self.sobel_y, padding=1)
        target_edge_x = F.conv2d(target_gray, self.sobel_x, padding=1)
        target_edge_y = F.conv2d(target_gray, self.sobel_y, padding=1)
        
        # 数值稳定性：添加小的 epsilon 避免 sqrt(0)
        pred_edge = torch.sqrt(pred_edge_x ** 2 + pred_edge_y ** 2 + 1e-6)
        target_edge = torch.sqrt(target_edge_x ** 2 + target_edge_y ** 2 + 1e-6)
        
        return F.l1_loss(pred_edge, target_edge)

class CombinedLoss(nn.Module):
    """组合损失函数 - 推荐用于手写去除"""
    def __init__(self, lambda_char=1.0, lambda_perc=0.1, lambda_ssim=0.5, lambda_edge=0.5):
        super().__init__()
        self.charbonnier = CharbonnierLoss()
        self.perceptual = PerceptualLoss()
        self.ssim = SSIMLoss()
        self.edge = EdgeLoss()
        
        self.lambda_char = lambda_char
        self.lambda_perc = lambda_perc
        self.lambda_ssim = lambda_ssim
        self.lambda_edge = lambda_edge
    
    def forward(self, pred, target):
        # 确保输入在合理范围内，避免异常值影响感知损失计算
        pred_clamped = torch.clamp(pred, 0, 1)
        
        loss_char = self.charbonnier(pred, target)
        loss_perc = self.perceptual(pred_clamped, target)
        loss_ssim = self.ssim(pred_clamped, target)
        loss_edge = self.edge(pred_clamped, target)
        
        # NaN 检查
        if torch.isnan(loss_char) or torch.isnan(loss_perc) or torch.isnan(loss_ssim) or torch.isnan(loss_edge):
            print(f"NaN detected! char={loss_char.item()}, perc={loss_perc.item()}, ssim={loss_ssim.item()}, edge={loss_edge.item()}")
        
        total_loss = (self.lambda_char * loss_char + 
                     self.lambda_perc * loss_perc + 
                     self.lambda_ssim * loss_ssim + 
                     self.lambda_edge * loss_edge)
        
        return total_loss, {
            'char': loss_char.item(),
            'perc': loss_perc.item(),
            'ssim': loss_ssim.item(),
            'edge': loss_edge.item()
        }

def train_model(args):
    # 1. Setup Device
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f'Using device: {device}')

    # 2. Prepare Data
    full_dataset = HandwriteDataset(
        img_dir=args.train_img_dir,
        label_dir=args.train_label_dir,
        size=(args.img_size, args.img_size),
        is_train=True
    )

    # Split into train and validation (90% train, 10% val)
    n_val = int(len(full_dataset) * 0.1)
    n_train = len(full_dataset) - n_val
    train_set, val_set = random_split(full_dataset, [n_train, n_val], generator=torch.Generator().manual_seed(42))

    train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, num_workers=4, pin_memory=True)
    val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, num_workers=4, pin_memory=True)

    print(f'Training images: {n_train}, Validation images: {n_val}')

    # 3. Initialize Model
    model = UNet(n_channels=3, n_classes=3, bilinear=True)
    model.to(device)
    
    # Use DataParallel for multi-GPU training
    if torch.cuda.device_count() > 1:
        print(f'Using {torch.cuda.device_count()} GPUs')
        model = nn.DataParallel(model)
    else:
        print('Using single GPU or CPU')

    # 4. Optimizer and Loss
    optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=1e-5)
    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5)
    
    # 根据参数选择损失函数
    if args.loss == 'combined':
        # 预设权重配置
        presets = {
            'conservative': {'char': 5.0, 'perc': 0.05, 'ssim': 0.3, 'edge': 0.3},
            'balanced': {'char': 10.0, 'perc': 0.05, 'ssim': 0.5, 'edge': 0.5},
            'aggressive': {'char': 20.0, 'perc': 0.1, 'ssim': 0.8, 'edge': 0.8}
        }
        weights = presets[args.loss_preset]
        criterion = CombinedLoss(
            lambda_char=weights['char'],
            lambda_perc=weights['perc'],
            lambda_ssim=weights['ssim'],
            lambda_edge=weights['edge']
        )
        print(f'Using Combined Loss ({args.loss_preset} preset): '
              f'Char={weights["char"]}, Perc={weights["perc"]}, SSIM={weights["ssim"]}, Edge={weights["edge"]}')
    elif args.loss == 'charbonnier':
        criterion = CharbonnierLoss()
        print('Using Charbonnier Loss')
    elif args.loss == 'mse':
        criterion = nn.MSELoss()
        print('Using MSE Loss')
    elif args.loss == 'l1':
        criterion = nn.L1Loss()
        print('Using L1 Loss') 

    # 5. Logging
    writer = SummaryWriter(log_dir=args.log_dir)
    os.makedirs(args.checkpoint_dir, exist_ok=True)

    # 6. Training Loop
    global_step = 0
    best_val_loss = float('inf')
    scaler = torch.amp.GradScaler('cuda')

    for epoch in range(args.epochs):
        model.train()
        epoch_loss = 0
        with tqdm(total=n_train, desc=f'Epoch {epoch + 1}/{args.epochs}', unit='img') as pbar:
            for batch in train_loader:
                images, true_masks = batch
                images = images.to(device)
                true_masks = true_masks.to(device)

                # Forward pass
                with torch.amp.autocast('cuda'):
                    masks_pred = model(images)
                    # 根据损失函数类型处理返回值
                    if args.loss == 'combined':
                        loss, loss_dict = criterion(masks_pred, true_masks)
                    else:
                        loss = criterion(masks_pred, true_masks)
                        loss_dict = {}

                # Backward and optimize
                optimizer.zero_grad()
                scaler.scale(loss).backward()
                
                # 梯度裁剪防止梯度爆炸
                scaler.unscale_(optimizer)
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
                
                # NaN 检测
                if torch.isnan(loss):
                    print(f"\nWarning: NaN loss detected at step {global_step}, skipping batch")
                    optimizer.zero_grad()
                    continue
                
                scaler.step(optimizer)
                scaler.update()

                # 显示进度
                postfix_dict = {'loss': f'{loss.item():.4f}', 'lr': f'{optimizer.param_groups[0]["lr"]:.2e}'}
                if loss_dict:
                    postfix_dict.update({k: f'{v:.4f}' for k, v in loss_dict.items()})
                
                # 监控输出范围，检测模型崩溃
                pred_min, pred_max = masks_pred.min().item(), masks_pred.max().item()
                if global_step % 100 == 0:
                    postfix_dict['pred_range'] = f'[{pred_min:.2f},{pred_max:.2f}]'
                
                pbar.set_postfix(**postfix_dict)
                pbar.update(images.shape[0])
                
                global_step += 1
                epoch_loss += loss.item()
                writer.add_scalar('Loss/train', loss.item(), global_step)
                # 记录各个损失分量（仅在使用 combined loss 时）
                if loss_dict:
                    for key, val in loss_dict.items():
                        writer.add_scalar(f'Loss/train_{key}', val, global_step)

        # Validation
        val_loss = evaluate(model, val_loader, criterion, device)
        writer.add_scalar('Loss/val', val_loss, global_step)
        print(f'Validation Loss: {val_loss:.4f}')

        # Update Scheduler
        scheduler.step(val_loss)

        # Save sample images to TensorBoard
        writer.add_images('images', images[:4], global_step)
        writer.add_images('masks/true', true_masks[:4], global_step)
        writer.add_images('masks/pred', masks_pred[:4], global_step)

        # Save Checkpoint
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            # Save model (handle DataParallel wrapper)
            model_to_save = model.module if hasattr(model, 'module') else model
            torch.save(model_to_save.state_dict(), os.path.join(args.checkpoint_dir, 'best_model.pth'))
            print(f'Saved best model with val loss: {val_loss:.4f}')
        
        # Save latest
        model_to_save = model.module if hasattr(model, 'module') else model
        torch.save(model_to_save.state_dict(), os.path.join(args.checkpoint_dir, 'latest_model.pth'))

    writer.close()

def evaluate(model, dataloader, criterion, device):
    model.eval()
    val_loss = 0
    with torch.no_grad():
        for batch in dataloader:
            images, true_masks = batch
            images = images.to(device)
            true_masks = true_masks.to(device)

            masks_pred = model(images)
            # 处理不同损失函数的返回值格式
            if isinstance(criterion, CombinedLoss):
                loss, _ = criterion(masks_pred, true_masks)
            else:
                loss = criterion(masks_pred, true_masks)
            val_loss += loss.item()
    
    return val_loss / len(dataloader)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train U-Net for Handwriting Removal')
    parser.add_argument('--train_img_dir', type=str, default='data/train', help='Path to training images')
    parser.add_argument('--train_label_dir', type=str, default='data/train_label', help='Path to training labels')
    parser.add_argument('--epochs', type=int, default=200, help='Number of epochs')
    parser.add_argument('--batch_size', type=int, default=4, help='Batch size')
    parser.add_argument('--lr', type=float, default=5e-5, help='Learning rate (default: 5e-5, conservative for combined loss)')
    parser.add_argument('--img_size', type=int, default=1024, help='Image size (must be divisible by 16)')
    parser.add_argument('--checkpoint_dir', type=str, default='checkpoints', help='Directory to save checkpoints')
    parser.add_argument('--log_dir', type=str, default='runs', help='Directory for TensorBoard logs')
    
    # Loss function options
    parser.add_argument('--loss', type=str, default='combined', 
                        choices=['charbonnier', 'mse', 'l1', 'combined'],
                        help='Loss function (default: combined)')
    parser.add_argument('--loss-preset', type=str, default='balanced',
                        choices=['conservative', 'balanced', 'aggressive'],
                        help='Preset weights for combined loss: conservative (stable), balanced (default), aggressive (sharp)')
    
    args = parser.parse_args()
    train_model(args)