04 - 训练目标推导¶

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

class DDPMTrainer:
    """
    DDPM训练器
    """
    def __init__(self, model, timesteps=1000, beta_schedule='linear', lr=1e-4):
        """
        参数:
            model: 噪声预测网络
            timesteps: 扩散步数
            beta_schedule: 噪声调度策略
            lr: 学习率
        """
        self.model = model
        self.timesteps = timesteps

        # 设置beta调度
        self.betas = self._get_beta_schedule(beta_schedule, timesteps)
        self.alphas = 1.0 - self.betas
        self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
        self.sqrt_alphas_cumprod = torch.sqrt(self.alphas_cumprod)
        self.sqrt_one_minus_alphas_cumprod = torch.sqrt(1.0 - self.alphas_cumprod)

        # 优化器
        self.optimizer = optim.Adam(model.parameters(), lr=lr)

        # 记录损失
        self.losses = []

    def _get_beta_schedule(self, schedule, timesteps):
        """获取beta调度"""
        if schedule == 'linear':
            return torch.linspace(0.0001, 0.02, timesteps)
        elif schedule == 'cosine':
            # 余弦调度实现
            s = 0.008
            steps = timesteps + 1
            x = torch.linspace(0, timesteps, steps)
            alphas_cumprod = torch.cos(((x / timesteps) + s) / (1 + s) * torch.pi / 2) ** 2
            alphas_cumprod = alphas_cumprod / alphas_cumprod[0]
            betas = 1 - (alphas_cumprod[1:] / alphas_cumprod[:-1])
            return torch.clip(betas, 0.0001, 0.9999)
        else:
            raise ValueError(f"Unknown schedule: {schedule}")

    def q_sample(self, x_0, t, noise=None):
        """
        前向扩散过程: 从 q(x_t | x_0) 采样

        参数:
            x_0: 原始图像 [B, C, H, W]
            t: 时间步 [B]
            noise: 可选的噪声

        返回:
            x_t: 加噪后的图像
            noise: 使用的噪声
        """
        if noise is None:
            noise = torch.randn_like(x_0)

        # 获取对应时间步的值
        sqrt_alpha_cumprod_t = self.sqrt_alphas_cumprod[t].view(-1, 1, 1, 1)  # 重塑张量形状
        sqrt_one_minus_alpha_cumprod_t = self.sqrt_one_minus_alphas_cumprod[t].view(-1, 1, 1, 1)

        # 重参数化
        x_t = sqrt_alpha_cumprod_t * x_0 + sqrt_one_minus_alpha_cumprod_t * noise

        return x_t, noise

    def training_step(self, x_0):
        """
        单步训练

        参数:
            x_0: 原始图像批次 [B, C, H, W]

        返回:
            loss: 损失值
        """
        self.model.train()  # train()训练模式
        batch_size = x_0.shape[0]
        device = x_0.device

        # 将调度参数移到设备
        self.sqrt_alphas_cumprod = self.sqrt_alphas_cumprod.to(device)  # 移至GPU/CPU
        self.sqrt_one_minus_alphas_cumprod = self.sqrt_one_minus_alphas_cumprod.to(device)

        # 1. 随机选择时间步
        t = torch.randint(0, self.timesteps, (batch_size,), device=device).long()

        # 2. 采样噪声
        noise = torch.randn_like(x_0)

        # 3. 计算加噪图像
        x_t, _ = self.q_sample(x_0, t, noise)

        # 4. 预测噪声
        noise_pred = self.model(x_t, t)

        # 5. 计算损失 (MSE)
        loss = nn.functional.mse_loss(noise_pred, noise)

        # 6. 反向传播
        self.optimizer.zero_grad()  # 清零梯度
        loss.backward()  # 反向传播计算梯度

        # 梯度裁剪（防止爆炸）
        torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)

        self.optimizer.step()  # 更新参数

        return loss.item()  # 将单元素张量转为Python数值

    def train_epoch(self, dataloader):
        """
        训练一个epoch

        参数:
            dataloader: 数据加载器

        返回:
            avg_loss: 平均损失
        """
        epoch_losses = []

        for batch_idx, (x_0, _) in enumerate(dataloader):  # enumerate同时获取索引和元素
            # 移动数据到设备
            x_0 = x_0.to(next(self.model.parameters()).device)

            # 训练步骤
            loss = self.training_step(x_0)
            epoch_losses.append(loss)

            # 打印进度
            if batch_idx % 100 == 0:
                print(f"Batch {batch_idx}/{len(dataloader)}, Loss: {loss:.4f}")

        avg_loss = sum(epoch_losses) / len(epoch_losses)
        self.losses.extend(epoch_losses)

        return avg_loss

# 使用示例
def train_ddpm():
    """训练DDPM的完整示例"""

    # 假设有数据加载器
    # dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

    # 创建模型（这里使用简化版UNet）
    # model = UNet(in_channels=3, out_channels=3)

    # 创建训练器
    # trainer = DDPMTrainer(model, timesteps=1000)

    # 训练循环
    # for epoch in range(num_epochs):
    #     avg_loss = trainer.train_epoch(dataloader)
    #     print(f"Epoch {epoch}, Average Loss: {avg_loss:.4f}")

    pass

# 使用余弦退火
scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=num_epochs)

# 或使用预热
from torch.optim.lr_scheduler import LambdaLR

def warmup_scheduler(epoch, warmup_epochs=10):
    if epoch < warmup_epochs:
        return epoch / warmup_epochs
    return 1.0

scheduler = LambdaLR(optimizer, lr_lambda=lambda epoch: warmup_scheduler(epoch))  # lambda匿名函数

class EMA:
    """指数移动平均"""
    def __init__(self, model, decay=0.9999):
        self.model = model
        self.decay = decay
        self.shadow = {}
        self.backup = {}

        # 初始化shadow参数
        for name, param in model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = param.data.clone()

    def update(self):
        """更新shadow参数"""
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.shadow[name] = self.decay * self.shadow[name] + (1 - self.decay) * param.data

    def apply_shadow(self):
        """使用shadow参数"""
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.backup[name] = param.data
                param.data = self.shadow[name]

    def restore(self):
        """恢复原始参数"""
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                param.data = self.backup[name]
        self.backup = {}

# 使用
ema = EMA(model)

# 训练循环中
for epoch in range(num_epochs):
    for x_0 in dataloader:
        loss = trainer.training_step(x_0)
        ema.update()

    # 采样时使用EMA参数
    ema.apply_shadow()
    # sample(...)
    ema.restore()

from torch.amp import autocast, GradScaler

scaler = GradScaler()

def training_step_mixed_precision(self, x_0):
    """混合精度训练步骤"""
    self.model.train()
    batch_size = x_0.shape[0]
    device = x_0.device

    t = torch.randint(0, self.timesteps, (batch_size,), device=device).long()
    noise = torch.randn_like(x_0)
    x_t, _ = self.q_sample(x_0, t, noise)

    # 使用自动混合精度
    with autocast('cuda'):
        noise_pred = self.model(x_t, t)
        loss = nn.functional.mse_loss(noise_pred, noise)

    # 缩放梯度
    self.optimizer.zero_grad()
    scaler.scale(loss).backward()
    scaler.unscale_(self.optimizer)
    torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
    scaler.step(self.optimizer)
    scaler.update()

    return loss.item()