扩散模型(Diffusion Model)原理与图像生成实战
扩散模型(Diffusion Model)原理,涵盖前向扩散添加噪声与反向扩散去噪过程。通过数学推导阐述 DDPM 训练目标,利用 PyTorch 构建 UNet 噪声预测网络。实战环节演示基于 MNIST 数据集的手写数字生成,包含数据加载、模型训练、损失计算及采样生成全流程。文末总结扩散模型变体及在图像生成、修复等领域的应用。

扩散模型(Diffusion Model)原理,涵盖前向扩散添加噪声与反向扩散去噪过程。通过数学推导阐述 DDPM 训练目标,利用 PyTorch 构建 UNet 噪声预测网络。实战环节演示基于 MNIST 数据集的手写数字生成,包含数据加载、模型训练、损失计算及采样生成全流程。文末总结扩散模型变体及在图像生成、修复等领域的应用。

💡 学习目标:掌握扩散模型的核心原理、前向扩散与反向扩散过程,以及基于扩散模型的图像生成任务实战流程。 💡 学习重点:理解扩散模型的噪声添加与噪声消除机制,学会使用 PyTorch 搭建 DDPM 模型,完成手写数字图像生成任务。
💡 传统的生成模型(如 GAN)存在训练不稳定、模式崩溃等问题。扩散模型作为一种基于概率的生成模型,通过逐步添加噪声和逐步去除噪声的双向过程,实现了更稳定的训练和更高质量的生成效果。 扩散模型的灵感来源于非平衡热力学,它的核心是将复杂的生成问题拆解为多个简单的马尔可夫链步骤。在图像生成、文本生成、语音合成等领域,扩散模型的表现已经超越了传统生成模型。
💡 扩散模型包含两个核心过程:前向扩散过程和反向扩散过程。
整个过程遵循马尔可夫链的假设,即每一步的状态只与前一步有关。
💡 前向扩散过程是一个固定的、非训练的过程。它的目标是通过逐步添加噪声,将真实图像 $x_0$ 转换为随机噪声 $x_T$。
前向扩散过程的每一步,都会按照以下公式向图像中添加噪声: $x_t = \sqrt{\alpha_t} x_{t-1} + \sqrt{1 - \alpha_t} \epsilon_t$ 其中:
为了计算方便,通常会定义累计乘积系数: $\bar{\alpha}t = \prod{i=1}^t \alpha_i$ 通过累计系数,可以直接从 $x_0$ 计算出任意步的 $x_t$: $x_t = \sqrt{\bar{\alpha}_t} x_0 + \sqrt{1 - \bar{\alpha}_t} \epsilon$
⚠️ 注意:前向扩散的步数 T 是一个超参数。T 越大,前向扩散越充分,反向扩散的效果越好,但训练和生成的时间也会越长。
import torch
import numpy as np
import matplotlib.pyplot as plt
# 定义扩散过程的超参数
T = 1000 # 扩散步数
beta_start = 0.0001 # 初始噪声系数
beta_end = 0.02 # 最终噪声系数
# 生成线性变化的 beta 序列
beta = torch.linspace(beta_start, beta_end, T)
alpha = 1 - beta
alpha_bar = torch.cumprod(alpha, dim=0) # 累计乘积
# 前向扩散函数:从 x0 生成 xt
def forward_diffusion(x0, t, device):
"""
x0: 原始图像 (batch_size, channels, height, width)
t: 扩散步数 (batch_size,)
"""
# 生成高斯噪声
eps = torch.randn_like(x0).to(device)
# 获取累计系数
alpha_bar_t = alpha_bar[t].reshape(-1, 1, 1, 1).to(device)
# 计算 xt
xt = torch.sqrt(alpha_bar_t) * x0 + torch.sqrt(1 - alpha_bar_t) * eps
return xt, eps
# 测试前向扩散过程
# 加载 MNIST 数据集的一张图像作为示例
from torchvision.datasets import MNIST
from torchvision.transforms import ToTensor
dataset = MNIST(root='./data', train=True, download=True, transform=ToTensor())
x0, _ = dataset[0]
x0 = x0.unsqueeze(0) # 增加 batch 维度 (1, 1, 28, 28)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 可视化不同步数的扩散效果
plt.figure(figsize=(15, 3))
for i, t in enumerate([0, 100, 200, 500, 800, 999]):
xt, _ = forward_diffusion(x0, torch.tensor([t]), device)
xt = xt.squeeze().cpu().detach().numpy()
plt.subplot(1, 6, i + 1)
plt.imshow(xt, cmap='gray')
plt.title(f't={t}')
plt.axis('off')
plt.show()
💡 反向扩散过程是前向扩散的逆过程。它的目标是训练一个神经网络 $\epsilon_\theta$,从 $x_t$ 中预测出添加的噪声 $\epsilon$,然后逐步去除噪声,还原出 $x_0$。
反向扩散的核心公式为: $p_\theta(x_{t-1}|x_t) = \mathcal{N}(x_{t-1}; \mu_\theta(x_t, t), \Sigma_\theta(x_t, t))$ 其中,均值 $\mu_\theta$ 由神经网络预测的噪声计算得到,方差 $\Sigma_\theta$ 通常设置为固定值。
训练的目标是最小化预测噪声和真实噪声之间的均方误差: $L(\theta) = \mathbb{E}{x_0, \epsilon, t}[|\epsilon - \epsilon\theta(x_t, t)|^2]$
💡 噪声预测网络是一个卷积神经网络。它的输入是 $x_t$ 和步数 t 的嵌入,输出是预测的噪声 $\epsilon_\theta$。
import torch.nn as nn
import torch.nn.functional as F
# 定义位置嵌入层:将步数 t 转换为高维向量
class PositionalEncoding(nn.Module):
def __init__(self, dim):
super().__init__()
self.dim = dim
def forward(self, t):
# t: (batch_size,)
device = t.device
half_dim = self.dim // 2
emb = np.log(10000) / (half_dim - 1)
emb = torch.exp(torch.arange(half_dim, device=device) * -emb)
emb = t[:, None] * emb[None, :]
emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
return emb
# 定义残差块
class ResidualBlock(nn.Module):
def __init__(self, in_channels, out_channels, time_dim):
super().__init__()
self.conv1 = nn.Conv2d(in_channels, out_channels, 3, padding=1)
self.bn1 = nn.BatchNorm2d(out_channels)
self.conv2 = nn.Conv2d(out_channels, out_channels, 3, padding=1)
self.bn2 = nn.BatchNorm2d(out_channels)
self.time_mlp = nn.Linear(time_dim, out_channels)
self.skip = nn.Conv2d(in_channels, out_channels, 1) if in_channels != out_channels else nn.Identity()
():
h = F.relu(.bn1(.conv1(x)))
h += .time_mlp(t)[:, :, , ]
h = F.relu(.bn2(.conv2(h)))
h + .skip(x)
(nn.Module):
():
().__init__()
.time_dim = time_dim
.pos_encoding = PositionalEncoding(time_dim)
.down1 = ResidualBlock(in_channels, , time_dim)
.down2 = ResidualBlock(, , time_dim)
.down3 = ResidualBlock(, , time_dim)
.pool = nn.MaxPool2d()
.bottleneck = ResidualBlock(, , time_dim)
.up1 = nn.ConvTranspose2d(, , , stride=)
.res_up1 = ResidualBlock(, , time_dim)
.up2 = nn.ConvTranspose2d(, , , stride=)
.res_up2 = ResidualBlock(, , time_dim)
.out = nn.Conv2d(, out_channels, )
():
t = .pos_encoding(t)
h1 = .down1(x, t)
h2 = .down2(.pool(h1), t)
h3 = .down3(.pool(h2), t)
bottleneck = .bottleneck(.pool(h3), t)
up1 = .up1(bottleneck)
up1 = torch.cat([up1, h3], dim=)
up1 = .res_up1(up1, t)
up2 = .up2(up1)
up2 = torch.cat([up2, h2], dim=)
up2 = .res_up2(up2, t)
up3 = .up2(up2)
up3 = torch.cat([up3, h1], dim=)
up3 = .res_up2(up3, t)
.out(up3)
model = UNet().to(device)
(model)
from torch.utils.data import DataLoader
from torch.optim import Adam
# 加载数据集
dataset = MNIST(root='./data', train=True, download=True, transform=ToTensor())
dataloader = DataLoader(dataset, batch_size=128, shuffle=True)
# 定义优化器
optimizer = Adam(model.parameters(), lr=1e-4)
criterion = nn.MSELoss()
# 训练函数
def train_epoch(model, dataloader, optimizer, criterion, device):
model.train()
total_loss = 0
for x0, _ in dataloader:
x0 = x0.to(device)
batch_size = x0.shape[0]
# 随机采样步数 t
t = torch.randint(0, T, (batch_size,), device=device)
# 前向扩散生成 xt 和真实噪声
xt, eps_true = forward_diffusion(x0, t, device)
# 模型预测噪声
eps_pred = model(xt, t)
# 计算损失
loss = criterion(eps_pred, eps_true)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
return total_loss / len(dataloader)
# 开始训练
epochs = 50
for epoch in range(epochs):
loss = train_epoch(model, dataloader, optimizer, criterion, device)
print(f"Epoch [{epoch+1}/{epochs}], Loss: {loss:.4f}")
# 保存模型
torch.save(model.state_dict(), 'ddpm_mnist.pth')
💡 训练完成后,我们可以从随机噪声出发,通过反向扩散过程逐步去除噪声,生成新的图像。
# 反向扩散采样函数
def sample(model, batch_size, device):
model.eval()
# 从随机噪声开始
xt = torch.randn((batch_size, 1, 28, 28)).to(device)
with torch.no_grad():
for t in reversed(range(1, T)):
# 生成当前步数的 tensor
t_tensor = torch.tensor([t], device=device).repeat(batch_size)
# 预测噪声
eps_pred = model(xt, t_tensor)
# 获取系数
alpha_t = alpha[t].to(device)
alpha_bar_t = alpha_bar[t].to(device)
alpha_bar_t_1 = alpha_bar[t-1].to(device)
beta_t = beta[t].to(device)
# 计算均值
mean = (1 / torch.sqrt(alpha_t)) * (xt - (beta_t / torch.sqrt(1 - alpha_bar_t)) * eps_pred)
# 计算方差
if t == 1:
variance = 0
else:
variance = beta_t
# 添加噪声
z = torch.randn_like(xt).to(device) if t > 1 else torch.zeros_like(xt).to(device)
xt = mean + torch.sqrt(variance) * z
# 归一化到 [0,1]
xt = torch.clamp(xt, 0, 1)
return xt
# 生成图像
model.load_state_dict(torch.load('ddpm_mnist.pth'))
generated_images = sample(model, batch_size=16, device=device)
# 可视化生成结果
plt.figure(figsize=(8, 8))
i ():
img = generated_images[i].squeeze().cpu().detach().numpy()
plt.subplot(, , i + )
plt.imshow(img, cmap=)
plt.axis()
plt.show()
💡 技巧 1:使用余弦噪声调度。将 beta 的变化从线性改为余弦,可以提升生成图像的质量。 💡 技巧 2:添加分类引导。在生成过程中加入类别信息,实现指定类别的图像生成。 💡 技巧 3:使用更大的网络架构。如 UNet++、Attention UNet 等,可以提升模型的特征捕捉能力。
✅ 扩散模型通过前向扩散和反向扩散两个过程,实现了稳定的生成模型训练,解决了 GAN 的训练不稳定问题。 ✅ 前向扩散是固定的噪声添加过程,反向扩散通过训练神经网络预测噪声,逐步还原真实数据。 ✅ 基于 DDPM 模型,可以实现手写数字等简单图像的生成,优化噪声调度和网络架构可提升生成效果。 ✅ 扩散模型的变体 Stable Diffusion 等已成为主流的图像生成工具,广泛应用于各类创意设计场景。

微信公众号「极客日志」,在微信中扫描左侧二维码关注。展示文案:极客日志 zeeklog
使用加密算法(如AES、TripleDES、Rabbit或RC4)加密和解密文本明文。 在线工具,加密/解密文本在线工具,online
生成新的随机RSA私钥和公钥pem证书。 在线工具,RSA密钥对生成器在线工具,online
基于 Mermaid.js 实时预览流程图、时序图等图表,支持源码编辑与即时渲染。 在线工具,Mermaid 预览与可视化编辑在线工具,online
解析常见 curl 参数并生成 fetch、axios、PHP curl 或 Python requests 示例代码。 在线工具,curl 转代码在线工具,online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具,Base64 字符串编码/解码在线工具,online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具,Base64 文件转换器在线工具,online