当前位置：首页 > news >正文

【第五章:计算机视觉-项目实战之生成式算法实战：扩散模型】3.生成式算法实战：扩散模型-(1)从零开始训练自己的扩散模型

news 2025/10/11 6:16:44

第五章：计算机视觉-项目实战之生成式算法实战：扩散模型

第三部分：生成式算法实战：扩散模型

第一节：从零开始训练自己的扩散模型

一、前言

在前两部分中，我们系统地学习了扩散模型（Diffusion Model）的理论基础与代表模型（如 DDPM、Stable Diffusion）。
本节我们将进入实战环节——带你 从零开始实现并训练一个属于自己的扩散模型！

我们的目标是：

构建一个 最小可行（Minimal Working） 的扩散模型框架，
使用 MNIST 或 CIFAR-10 等轻量级数据集完成训练与图像生成。

二、扩散模型回顾

扩散模型的核心思想非常优雅：

通过一个逐步加噪（Diffusion）与逐步去噪（Denoising）的过程，
实现从纯噪声到图像的生成。

简要流程如下：

1.前向扩散过程（Forward Process）：
逐步往真实图像中加入噪声，使其变得越来越模糊，直到变成纯随机噪声。

2.反向扩散过程（Reverse Process）：
训练一个神经网络（通常是 U-Net），学习如何一步步从噪声中“还原”出原始图像。

三、实验准备

1.环境依赖

pip install torch torchvision matplotlib tqdm

2.数据集准备

我们使用经典的 MNIST 手写数字数据集 作为示例。

from torchvision import datasets, transforms
from torch.utils.data import DataLoadertransform = transforms.Compose([transforms.ToTensor(),transforms.Normalize((0.5,), (0.5,))
])dataset = datasets.MNIST(root="./data", train=True, download=True, transform=transform)
dataloader = DataLoader(dataset, batch_size=128, shuffle=True)

四、核心数学推导（简化版）

扩散模型的加噪公式如下：

$x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t}\epsilon$

其中：

$x_0$ ：原始图像
$\epsilon$ ：高斯噪声
$\bar{\alpha}t = \prod{s=1}^{t}(1 - \beta_s)$ ：噪声累积系数

训练目标是让网络 $\epsilon_\theta(x_t, t)$ 预测噪声 $\epsilon$ 。
损失函数：

$L = \mathbb{E}{x_0, t, \epsilon} [||\epsilon - \epsilon\theta(x_t, t)||^2]$

五、模型架构实现（简化 U-Net）

我们使用一个小型 U-Net 结构来预测噪声：

import torch
import torch.nn as nn
import torch.nn.functional as Fclass SimpleUNet(nn.Module):def __init__(self, channels=1):super().__init__()self.down1 = nn.Conv2d(channels, 64, 3, padding=1)self.down2 = nn.Conv2d(64, 128, 3, padding=1)self.up1 = nn.ConvTranspose2d(128, 64, 3, padding=1)self.out = nn.Conv2d(64, channels, 3, padding=1)def forward(self, x, t):h1 = F.relu(self.down1(x))h2 = F.relu(self.down2(h1))h3 = F.relu(self.up1(h2))out = self.out(h3)return out

六、扩散过程实现

我们定义噪声调度器（Noise Scheduler）：

import torchclass Diffusion:def __init__(self, timesteps=300, beta_start=1e-4, beta_end=0.02):self.timesteps = timestepsself.betas = torch.linspace(beta_start, beta_end, timesteps)self.alphas = 1. - self.betasself.alphas_cumprod = torch.cumprod(self.alphas, dim=0)def q_sample(self, x0, t, noise=None):if noise is None:noise = torch.randn_like(x0)sqrt_alphas = torch.sqrt(self.alphas_cumprod[t])[:, None, None, None]sqrt_one_minus = torch.sqrt(1 - self.alphas_cumprod[t])[:, None, None, None]return sqrt_alphas * x0 + sqrt_one_minus * noise

七、训练流程

import torch.optim as optim
from tqdm import tqdmdevice = "cuda" if torch.cuda.is_available() else "cpu"
model = SimpleUNet().to(device)
diffusion = Diffusion()
optimizer = optim.Adam(model.parameters(), lr=1e-4)epochs = 10
for epoch in range(epochs):for images, _ in tqdm(dataloader):images = images.to(device)t = torch.randint(0, diffusion.timesteps, (images.size(0),), device=device)noise = torch.randn_like(images)x_t = diffusion.q_sample(images, t, noise)noise_pred = model(x_t, t)loss = F.mse_loss(noise_pred, noise)optimizer.zero_grad()loss.backward()optimizer.step()print(f"Epoch {epoch+1}: Loss={loss.item():.4f}")

八、图像生成（采样过程）

训练完成后，我们可以从纯噪声生成新图像：

@torch.no_grad()
def sample(model, diffusion, size=16):x = torch.randn(size, 1, 28, 28).to(device)for t in reversed(range(diffusion.timesteps)):noise_pred = model(x, torch.tensor([t]*size).to(device))beta_t = diffusion.betas[t]alpha_t = diffusion.alphas[t]alpha_hat = diffusion.alphas_cumprod[t]if t > 0:noise = torch.randn_like(x)else:noise = 0x = (1 / torch.sqrt(alpha_t)) * (x - ((1 - alpha_t) / torch.sqrt(1 - alpha_hat)) * noise_pred) + torch.sqrt(beta_t) * noisereturn x# 生成样例
generated = sample(model, diffusion)

九、结果可视化

import matplotlib.pyplot as pltplt.figure(figsize=(6,6))
for i in range(16):plt.subplot(4,4,i+1)plt.imshow(generated[i].cpu().squeeze(), cmap="gray")plt.axis("off")
plt.show()

你会看到模型从随机噪声中逐渐生成出手写数字的轮廓！
虽然效果不及 Stable Diffusion，但这是一个完整的「自制扩散模型」训练闭环！

十、性能优化与扩展方向

优化点	说明
更深的 U-Net	提升模型容量与表达能力
位置嵌入（Positional Embedding）	增强模型对时间步 t 的感知能力
训练调度	使用 cosine 或 linear beta schedule 改善稳定性
条件生成	给模型加上标签条件，实现“可控生成”
高分辨率扩展	使用 CIFAR-10 / CelebA-HQ 数据集进行扩展实验