人工智能:扩散模型(Diffusion Model)原理与图像生成实战

AI2天前发布 beixibaobao
1 0 0

人工智能:扩散模型(Diffusion Model)原理与图像生成实战

在这里插入图片描述

1.1 本章学习目标与重点

💡 学习目标:掌握扩散模型的核心原理、前向扩散与反向扩散过程,以及基于扩散模型的图像生成任务实战流程。
💡 学习重点:理解扩散模型的噪声添加与噪声消除机制,学会使用 PyTorch 搭建 DDPM 模型,完成手写数字图像生成任务。

1.2 扩散模型的核心思想

1.2.1 为什么需要扩散模型

💡 传统的生成模型(如 GAN)存在训练不稳定模式崩溃等问题。扩散模型作为一种基于概率的生成模型,通过逐步添加噪声和逐步去除噪声的双向过程,实现了更稳定的训练和更高质量的生成效果。
扩散模型的灵感来源于非平衡热力学,它的核心是将复杂的生成问题拆解为多个简单的马尔可夫链步骤。在图像生成、文本生成、语音合成等领域,扩散模型的表现已经超越了传统生成模型。

1.2.2 扩散模型的基本框架

💡 扩散模型包含两个核心过程:前向扩散过程反向扩散过程

  1. 前向扩散过程:从真实数据出发,逐步向数据中添加高斯噪声。经过 T 步后,数据会变成完全随机的噪声。
  2. 反向扩散过程:从随机噪声出发,训练一个神经网络逐步去除噪声。经过 T 步后,噪声会还原为真实的数据分布。

整个过程遵循马尔可夫链的假设,即每一步的状态只与前一步有关。

1.3 前向扩散过程详解

💡 前向扩散过程是一个固定的、非训练的过程。它的目标是通过逐步添加噪声,将真实图像 x0x_0x0 转换为随机噪声 xTx_TxT

1.3.1 前向扩散的数学原理

前向扩散过程的每一步,都会按照以下公式向图像中添加噪声:
xt=αtxt−1+1−αtϵtx_t = sqrt{alpha_t} x_{t-1} + sqrt{1 – alpha_t} epsilon_txt=αtxt1+1αtϵt
其中:

  • xtx_txt 表示第 t 步添加噪声后的图像
  • αtalpha_tαt 是一个预先设定的噪声系数,满足 0<αt<10 < alpha_t < 10<αt<1
  • ϵtepsilon_tϵt 是服从标准正态分布的高斯噪声

为了计算方便,通常会定义累计乘积系数:
αtˉ=∏i=1tαibar{alpha_t} = prod_{i=1}^t alpha_iαtˉ=i=1tαi
通过累计系数,可以直接从 x0x_0x0 计算出任意步的 xtx_txt
xt=αtˉx0+1−αtˉϵx_t = sqrt{bar{alpha_t}} x_0 + sqrt{1 – bar{alpha_t}} epsilonxt=αtˉx0+1αtˉϵ

⚠️ 注意:前向扩散的步数 T 是一个超参数。T 越大,前向扩散越充分,反向扩散的效果越好,但训练和生成的时间也会越长。

1.3.2 前向扩散过程的代码实现

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()

1.4 反向扩散过程与模型训练

1.4.1 反向扩散的数学原理

💡 反向扩散过程是前向扩散的逆过程。它的目标是训练一个神经网络 ϵθepsilon_thetaϵθ,从 xtx_txt 中预测出添加的噪声 ϵepsilonϵ,然后逐步去除噪声,还原出 x0x_0x0

反向扩散的核心公式为:
pθ(xt−1∣xt)=N(xt−1;μθ(xt,t),Σθ(xt,t))p_theta(x_{t-1}|x_t) = mathcal{N}(x_{t-1}; mu_theta(x_t, t), Sigma_theta(x_t, t))pθ(xt1xt)=N(xt1;μθ(xt,t),Σθ(xt,t))
其中,均值 μθmu_thetaμθ 由神经网络预测的噪声计算得到,方差 ΣθSigma_thetaΣθ 通常设置为固定值。

训练的目标是最小化预测噪声和真实噪声之间的均方误差:
L(θ)=Ex0,ϵ,t[∥ϵ−ϵθ(xt,t)∥2]L(theta) = mathbb{E}_{x_0, epsilon, t}[|epsilon – epsilon_theta(x_t, t)|^2]L(θ)=Ex0,ϵ,t[ϵϵθ(xt,t)2]

1.4.2 构建噪声预测网络

💡 噪声预测网络是一个卷积神经网络。它的输入是 xtx_txt 和步数 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()
    def forward(self, x, t):
        # x: (batch_size, channels, h, w)
        # t: (batch_size, time_dim)
        h = F.relu(self.bn1(self.conv1(x)))
        # 添加时间嵌入
        h += self.time_mlp(t)[:, :, None, None]
        h = F.relu(self.bn2(self.conv2(h)))
        return h + self.skip(x)
# 定义噪声预测网络
class UNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1, time_dim=256):
        super().__init__()
        self.time_dim = time_dim
        self.pos_encoding = PositionalEncoding(time_dim)
        # 下采样路径
        self.down1 = ResidualBlock(in_channels, 64, time_dim)
        self.down2 = ResidualBlock(64, 128, time_dim)
        self.down3 = ResidualBlock(128, 256, time_dim)
        self.pool = nn.MaxPool2d(2)
        # 瓶颈层
        self.bottleneck = ResidualBlock(256, 256, time_dim)
        # 上采样路径
        self.up1 = nn.ConvTranspose2d(256, 128, 2, stride=2)
        self.res_up1 = ResidualBlock(256, 128, time_dim)
        self.up2 = nn.ConvTranspose2d(128, 64, 2, stride=2)
        self.res_up2 = ResidualBlock(128, 64, time_dim)
        # 输出层
        self.out = nn.Conv2d(64, out_channels, 1)
    def forward(self, x, t):
        # x: (batch_size, 1, 28, 28)
        # t: (batch_size,)
        # 位置编码
        t = self.pos_encoding(t)
        # 下采样
        h1 = self.down1(x, t)
        h2 = self.down2(self.pool(h1), t)
        h3 = self.down3(self.pool(h2), t)
        # 瓶颈层
        bottleneck = self.bottleneck(self.pool(h3), t)
        # 上采样
        up1 = self.up1(bottleneck)
        up1 = torch.cat([up1, h3], dim=1)
        up1 = self.res_up1(up1, t)
        up2 = self.up2(up1)
        up2 = torch.cat([up2, h2], dim=1)
        up2 = self.res_up2(up2, t)
        up3 = self.up2(up2)
        up3 = torch.cat([up3, h1], dim=1)
        up3 = self.res_up2(up3, t)
        return self.out(up3)
# 初始化模型
model = UNet().to(device)
print(model)

1.4.3 模型训练流程

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')

1.5 实战:基于DDPM的图像生成

1.5.1 反向扩散采样过程

💡 训练完成后,我们可以从随机噪声出发,通过反向扩散过程逐步去除噪声,生成新的图像。

# 反向扩散采样函数
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))
for i in range(16):
    img = generated_images[i].squeeze().cpu().detach().numpy()
    plt.subplot(4, 4, i+1)
    plt.imshow(img, cmap='gray')
    plt.axis('off')
plt.show()

1.5.2 模型优化技巧

💡 技巧1:使用余弦噪声调度。将 beta 的变化从线性改为余弦,可以提升生成图像的质量。
💡 技巧2:添加分类引导。在生成过程中加入类别信息,实现指定类别的图像生成。
💡 技巧3:使用更大的网络架构。如 UNet++、Attention UNet 等,可以提升模型的特征捕捉能力。

1.6 扩散模型的发展与应用

1.6.1 经典扩散模型变体

  • DDIM:Denoising Diffusion Implicit Models,通过简化反向扩散过程,大幅提升生成速度。
  • Stable Diffusion:基于潜在空间的扩散模型,在保持生成质量的同时,降低了计算资源消耗。
  • Latent Diffusion Models:在压缩的潜在空间中进行扩散,适用于高分辨率图像生成。

1.6.2 扩散模型的应用场景

  • 图像生成:生成艺术画作、人像、产品设计图等。
  • 图像修复:修复老照片、去除图像中的水印和瑕疵。
  • 文本到图像生成:根据文本描述生成对应的图像,如 Midjourney、DALL·E。
  • 语音合成:生成自然流畅的语音,提升语音合成的质量。

1.7 本章总结

✅ 扩散模型通过前向扩散和反向扩散两个过程,实现了稳定的生成模型训练,解决了 GAN 的训练不稳定问题。
✅ 前向扩散是固定的噪声添加过程,反向扩散通过训练神经网络预测噪声,逐步还原真实数据。
✅ 基于 DDPM 模型,可以实现手写数字等简单图像的生成,优化噪声调度和网络架构可提升生成效果。
✅ 扩散模型的变体 Stable Diffusion 等已成为主流的图像生成工具,广泛应用于各类创意设计场景。

© 版权声明

相关文章