1. 扩散模型基础原理拆解
扩散模型（Diffusion Model）作为当前生成式AI领域的核心技术之一，其核心思想源于物理学中的扩散过程。想象一滴墨水在清水中的扩散过程——从清晰的初始状态逐渐变成均匀的混合状态，而扩散模型则逆向学习这个过程。
1.1 前向扩散过程
前向扩散过程可以形式化为马尔可夫链，每一步都对当前图像x_t添加高斯噪声。数学表达为：
python复制q(x_t|x_{t-1}) = N(x_t; √(1-β_t)x_{t-1}, β_tI)
其中β_t是噪声调度参数，控制每一步的噪声强度。实际实现时通常采用线性调度：
python复制def linear_beta_schedule(timesteps):
beta_start = 0.0001
beta_end = 0.02
return torch.linspace(beta_start, beta_end, timesteps)
1.2 逆向去噪过程
逆向过程需要学习一个神经网络来预测添加的噪声：
python复制ε_θ(x_t, t) = model(x_t, t)
常用的UNet结构会包含：
- 下采样编码器（Encoder）
- 时间嵌入层（Time Embedding）
- 注意力机制模块
- 上采样解码器（Decoder）
关键技巧：在实现时使用GroupNorm替代BatchNorm，因为小批量训练时BN会导致性能不稳定。
2. PyTorch实现核心模块
2.1 噪声调度器实现
python复制class NoiseScheduler:
def __init__(self, timesteps=1000):
self.betas = linear_beta_schedule(timesteps)
self.alphas = 1. - self.betas
self.alphas_cumprod = torch.cumprod(self.alphas, dim=0)
def add_noise(self, x_start, noise, t):
sqrt_alpha_cumprod = torch.sqrt(self.alphas_cumprod[t])
sqrt_one_minus_alpha_cumprod = torch.sqrt(1. - self.alphas_cumprod[t])
return sqrt_alpha_cumprod * x_start + sqrt_one_minus_alpha_cumprod * noise
2.2 时间嵌入层
python复制class TimeEmbedding(nn.Module):
def __init__(self, dim):
super().__init__()
self.dim = dim
half_dim = dim // 2
emb = math.log(10000) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
self.register_buffer('emb', emb)
def forward(self, t):
emb = t.float() * self.emb
emb = torch.cat((emb.sin(), emb.cos()), dim=-1)
return emb
2.3 残差块设计
python复制class ResidualBlock(nn.Module):
def __init__(self, in_channels, out_channels, time_dim):
super().__init__()
self.conv1 = nn.Sequential(
nn.GroupNorm(32, in_channels),
nn.SiLU(),
nn.Conv2d(in_channels, out_channels, 3, padding=1)
)
self.time_mlp = nn.Linear(time_dim, out_channels)
self.conv2 = nn.Sequential(
nn.GroupNorm(32, out_channels),
nn.SiLU(),
nn.Conv2d(out_channels, out_channels, 3, padding=1)
)
self.res_conv = nn.Conv2d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()
def forward(self, x, t):
h = self.conv1(x)
h += self.time_mlp(t)[:, :, None, None]
h = self.conv2(h)
return h + self.res_conv(x)
3. 完整训练流程实现
3.1 数据准备与增强
建议使用torchvision的变换组合：
python复制transform = transforms.Compose([
transforms.Resize(256),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize([0.5], [0.5])
])
3.2 训练循环关键代码
python复制def train_loop(model, dataloader, optimizer, scheduler, device):
model.train()
for batch in dataloader:
optimizer.zero_grad()
x = batch.to(device)
noise = torch.randn_like(x)
t = torch.randint(0, scheduler.timesteps, (x.shape[0],)).to(device)
noisy_x = scheduler.add_noise(x, noise, t)
predicted_noise = model(noisy_x, t)
loss = F.mse_loss(predicted_noise, noise)
loss.backward()
optimizer.step()
3.3 采样生成图像
python复制@torch.no_grad()
def sample(model, scheduler, num_samples, size, device):
x = torch.randn((num_samples, 3, size, size)).to(device)
for t in reversed(range(scheduler.timesteps)):
model_input = torch.cat([x] * 2) # 用于分类器自由引导
noise_pred = model(model_input, t)
noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
noise_pred = noise_pred_uncond + 7.5 * (noise_pred_cond - noise_pred_uncond)
x = scheduler.step(x, noise_pred, t)
return x
4. 实战优化技巧
4.1 混合精度训练
python复制scaler = torch.cuda.amp.GradScaler()
with torch.cuda.amp.autocast():
predicted_noise = model(noisy_x, t)
loss = F.mse_loss(predicted_noise, noise)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
4.2 学习率调度策略
推荐使用WarmupCosineSchedule：
python复制def warmup_cosine_schedule(optimizer, warmup_steps, total_steps):
def lr_lambda(step):
if step < warmup_steps:
return float(step) / float(max(1, warmup_steps))
progress = float(step - warmup_steps) / float(max(1, total_steps - warmup_steps))
return 0.5 * (1. + math.cos(math.pi * progress))
return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
4.3 模型保存与恢复
python复制def save_checkpoint(model, optimizer, scheduler, epoch, path):
torch.save({
'model_state_dict': model.state_dict(),
'optimizer_state_dict': optimizer.state_dict(),
'scheduler_state_dict': scheduler.state_dict(),
'epoch': epoch,
}, path)
def load_checkpoint(path, model, optimizer=None, scheduler=None):
checkpoint = torch.load(path)
model.load_state_dict(checkpoint['model_state_dict'])
if optimizer:
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
if scheduler:
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
return checkpoint['epoch']
5. 常见问题排查指南
| 问题现象 | 可能原因 | 解决方案 |
|---|---|---|
| 生成图像全黑/全白 | 数值溢出 | 检查数据归一化范围，确保在[-1,1]之间 |
| 训练loss不下降 | 学习率设置不当 | 尝试warmup策略或降低学习率10倍 |
| GPU内存不足 | 批次过大/分辨率过高 | 减小batch_size或使用梯度累积 |
| 生成图像模糊 | 训练不充分 | 增加训练epoch或使用感知损失 |
| 采样速度慢 | timesteps过多 | 使用DDIM加速采样或减少timesteps |
重要提示：当使用fp16训练时，建议设置--gradient_clip_val=1.0避免梯度爆炸
6. 进阶改进方向
6.1 潜在扩散模型（LDM）
通过VAE将图像压缩到潜在空间：
python复制class VAE(nn.Module):
def __init__(self):
super().__init__()
self.encoder = nn.Sequential(
nn.Conv2d(3, 64, 3, stride=2, padding=1),
nn.ReLU(),
nn.Conv2d(64, 128, 3, stride=2, padding=1),
nn.ReLU()
)
self.decoder = nn.Sequential(
nn.ConvTranspose2d(128, 64, 3, stride=2, padding=1, output_padding=1),
nn.ReLU(),
nn.ConvTranspose2d(64, 3, 3, stride=2, padding=1, output_padding=1),
nn.Tanh()
)
6.2 分类器自由引导
在训练时随机丢弃文本条件：
python复制def forward(self, x, t, cond=None):
if cond is not None and self.training:
mask = torch.rand(cond.shape[0]) < 0.1
cond[mask] = 0
# 正常前向传播...
6.3 多模态扩散
扩展模型处理文本条件：
python复制class TextConditionedUNet(nn.Module):
def __init__(self, text_encoder):
super().__init__()
self.text_proj = nn.Linear(text_encoder.config.hidden_size, 768)
# 其余UNet结构...
def forward(self, x, t, text_embeddings):
text_emb = self.text_proj(text_embeddings)
# 将text_emb与时间嵌入结合...
在实际项目中，我发现扩散模型对超参数非常敏感。建议初始阶段使用小规模数据集（如CIFAR10）进行快速验证，待pipeline调通后再扩展到更大数据集。对于256x256以上分辨率，采用渐进式训练策略效果更好——先训练低分辨率模型，再逐步提高分辨率。