Python实现基于扩散模型（Diffusion Model）的生成式AI：采样与去噪过程
大家好，今天我们来深入探讨扩散模型，并用Python代码实现其核心的采样和去噪过程。扩散模型作为近年来生成式AI领域的一颗新星，以其独特的理论基础和出色的生成效果，受到了广泛的关注。
1. 扩散模型的核心思想
扩散模型的核心思想是将数据生成过程建模为一个马尔可夫链，该链包含两个过程：扩散过程（Forward Diffusion Process） 和 逆扩散过程（Reverse Diffusion Process）。
-
扩散过程： 从原始数据出发，逐步添加高斯噪声，直到数据完全变成噪声，失去原始数据的特征。这个过程通常是固定的，并且可以通过预定义的噪声时间表（noise schedule）来控制噪声添加的强度。
-
逆扩散过程： 从纯高斯噪声出发，逐步去除噪声，恢复出原始数据。这个过程是扩散模型的关键，它需要学习一个模型来预测每一步需要去除的噪声。
简单来说，扩散模型就像将一张照片逐渐模糊化，直到完全看不清，然后学习如何一步步地将模糊的照片恢复清晰。
2. 数学原理：前向扩散过程
前向扩散过程是一个马尔可夫过程，它从原始数据分布 x_0 ~ q(x) 开始，逐步添加高斯噪声。在每个时间步 t，我们向 x_{t-1} 添加噪声，得到 x_t。这个过程可以表示为：
q(x_t | x_{t-1}) = N(x_t; √(1 - β_t) * x_{t-1}, β_t * I)其中：
x_t是时间步t的数据。β_t是时间步t的噪声方差，通常是一个递增的序列，也称为噪声时间表。I是单位矩阵。N(μ, Σ)表示均值为μ，协方差矩阵为Σ的高斯分布。
利用马尔可夫性质，我们可以直接计算任意时间步 t 的 x_t，而无需逐步迭代。这个公式如下：
q(x_t | x_0) = N(x_t; √(α_t) * x_0, (1 - α_t) * I)其中：
α_t = ∏_{i=1}^{t} (1 - β_i)是一个累积的降噪系数。
3. 数学原理：反向扩散过程
反向扩散过程的目标是从纯高斯噪声 x_T ~ N(0, I) 开始，逐步去除噪声，恢复出原始数据 x_0。由于反向过程的分布是未知的，我们需要学习一个模型来近似它。通常，我们使用神经网络来学习这个模型。
反向过程可以表示为：
p_θ(x_{t-1} | x_t) = N(x_{t-1}; μ_θ(x_t, t), Σ_θ(x_t, t))其中：
μ_θ(x_t, t)是神经网络预测的均值。Σ_θ(x_t, t)是神经网络预测的方差。
扩散模型的目标是最小化KL散度，使得学习到的反向过程尽可能接近真实的反向过程。
4. 关键代码实现：噪声时间表
噪声时间表 β_t 的选择对扩散模型的性能至关重要。常见的噪声时间表包括线性、二次、余弦等。下面是一个线性噪声时间表的Python实现：
import torch
def linear_beta_schedule(timesteps, beta_start=0.0001, beta_end=0.02):
"""
生成线性噪声时间表.
Args:
timesteps: 总的时间步数.
beta_start: 噪声方差的起始值.
beta_end: 噪声方差的结束值.
Returns:
torch.Tensor: 噪声时间表.
"""
beta = torch.linspace(beta_start, beta_end, timesteps)
return beta
# 示例：生成1000个时间步的线性噪声时间表
timesteps = 1000
betas = linear_beta_schedule(timesteps)
print(betas)5. 关键代码实现：前向扩散过程
下面是前向扩散过程的Python实现：
def forward_diffusion_sample(x_0, t, betas):
"""
前向扩散过程采样.
Args:
x_0: 原始数据.
t: 时间步.
betas: 噪声时间表.
Returns:
torch.Tensor: 噪声数据 x_t.
torch.Tensor: 添加的噪声.
"""
sqrt_alpha_cumprod = torch.sqrt(torch.cumprod(1 - betas, dim=0))
sqrt_alpha_cumprod_t = sqrt_alpha_cumprod[t].view(-1, 1, 1, 1) # reshape to match x_0's shape
sqrt_one_minus_alpha_cumprod = torch.sqrt(1 - torch.cumprod(1 - betas, dim=0))
sqrt_one_minus_alpha_cumprod_t = sqrt_one_minus_alpha_cumprod[t].view(-1, 1, 1, 1)
noise = torch.randn_like(x_0)
x_t = sqrt_alpha_cumprod_t * x_0 + sqrt_one_minus_alpha_cumprod_t * noise
return x_t, noise
# 示例：对图像数据进行前向扩散
# 假设 x_0 是一个形状为 (1, 1, 28, 28) 的图像数据，表示一个单通道28x28的图像
# 注意：这个形状是示例，你需要根据你的实际图像数据调整
x_0 = torch.randn(1, 1, 28, 28) # 模拟一个图像数据
t = torch.tensor([500]) # 选择时间步
x_t, noise = forward_diffusion_sample(x_0, t, betas)
print("Shape of x_t:", x_t.shape)
print("Shape of noise:", noise.shape)解释：
-
sqrt_alpha_cumprod和sqrt_one_minus_alpha_cumprod的计算：torch.cumprod(1 - betas, dim=0)计算(1 - β_1) * (1 - β_2) * ... * (1 - β_t)，也就是α_t。 然后，我们计算α_t的平方根以及(1 - α_t)的平方根。 这些值在公式q(x_t | x_0) = N(x_t; √(α_t) * x_0, (1 - α_t) * I)中使用。 -
reshape操作：sqrt_alpha_cumprod_t = sqrt_alpha_cumprod[t].view(-1, 1, 1, 1)这行代码非常重要。sqrt_alpha_cumprod[t]会返回一个标量值（因为t是一个标量tensor）。 为了能够正确地与x_0进行广播相乘（element-wise multiplication），我们需要将这个标量值 reshape 成与x_0相同的维度，但除了第一个维度之外，其他维度的大小都为1。-1表示让PyTorch自动推断第一个维度的大小，这里它会是x_0的 batch size (在这个例子中是 1). 假设x_0的形状是(batch_size, channels, height, width)， 那么sqrt_alpha_cumprod_t的形状将会是(batch_size, 1, 1, 1)。 这样，当sqrt_alpha_cumprod_t乘以x_0时，它会沿着通道、高度和宽度维度进行广播，从而实现正确的缩放。sqrt_one_minus_alpha_cumprod_t的 reshape 操作也是同样的道理。 -
噪声的生成：
noise = torch.randn_like(x_0)生成一个与x_0形状相同的高斯噪声张量。 -
x_t的计算：x_t = sqrt_alpha_cumprod_t * x_0 + sqrt_one_minus_alpha_cumprod_t * noise使用公式x_t = √(α_t) * x_0 + √(1 - α_t) * noise计算x_t。
6. 关键代码实现：反向扩散过程（去噪）
反向扩散过程的核心是训练一个神经网络来预测噪声。 下面是一个简化的反向扩散过程的Python实现，假设我们已经训练好了一个名为 model 的神经网络，它可以根据 x_t 和时间步 t 预测噪声。
import torch.nn as nn
import torch.nn.functional as F
# 一个简单的U-Net模型，用于预测噪声
class SimpleUnet(nn.Module):
def __init__(self, in_channels=1, out_channels=1, time_dim=256):
super().__init__()
self.time_mlp = nn.Linear(time_dim, in_channels)
# Downsampling
self.conv1 = nn.Conv2d(in_channels, 32, kernel_size=3, padding=1)
self.conv2 = nn.Conv2d(32, 64, kernel_size=3, padding=1)
self.downsample1 = nn.MaxPool2d(2)
self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
self.conv4 = nn.Conv2d(128, 128, kernel_size=3, padding=1)
self.downsample2 = nn.MaxPool2d(2)
# Bottleneck
self.bottleneck_conv1 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
self.bottleneck_conv2 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
# Upsampling
self.upsample1 = nn.Upsample(scale_factor=2, mode='nearest')
self.conv_up1 = nn.Conv2d(128 + 128, 64, kernel_size=3, padding=1) # Increased input channels
self.conv_up2 = nn.Conv2d(64, 64, kernel_size=3, padding=1)
self.upsample2 = nn.Upsample(scale_factor=2, mode='nearest')
self.conv_up3 = nn.Conv2d(64 + 32, 32, kernel_size=3, padding=1) # Increased input channels
self.conv_up4 = nn.Conv2d(32, 32, kernel_size=3, padding=1)
# Output
self.output_conv = nn.Conv2d(32, out_channels, kernel_size=1)
def forward(self, x, t):
# Time embedding
t = self.time_mlp(t)
t = t.view(-1, x.shape[1], 1, 1) # Reshape to match image dimensions
x = x + t
# Downsampling
x1 = F.relu(self.conv1(x))
x2 = F.relu(self.conv2(x1))
x2_down = self.downsample1(x2)
x3 = F.relu(self.conv3(x2_down))
x4 = F.relu(self.conv4(x3))
x4_down = self.downsample2(x4)
# Bottleneck
x_bottleneck = F.relu(self.bottleneck_conv1(x4_down))
x_bottleneck = F.relu(self.bottleneck_conv2(x_bottleneck))
# Upsampling
x_up1 = self.upsample1(x_bottleneck)
x_up1 = torch.cat([x_up1, x4], dim=1) # Concatenate skip connection
x_up1 = F.relu(self.conv_up1(x_up1))
x_up1 = F.relu(self.conv_up2(x_up1))
x_up2 = self.upsample2(x_up1)
x_up2 = torch.cat([x_up2, x2], dim=1) # Concatenate skip connection
x_up2 = F.relu(self.conv_up3(x_up2))
x_up2 = F.relu(self.conv_up4(x_up2))
# Output
output = self.output_conv(x_up2)
return output
def reverse_diffusion(x_t, t, model, betas):
"""
反向扩散过程（去噪）.
Args:
x_t: 当前时间步的噪声数据.
t: 当前时间步.
model: 训练好的噪声预测模型.
betas: 噪声时间表.
Returns:
torch.Tensor: 去噪后的数据 x_{t-1}.
"""
# 将 t 转换为模型所需的格式 (例如，embedding)
t = t.float() / timesteps # 归一化时间步
# 简单的时间步编码, 可以替换为更复杂的编码方式
time_embedding = torch.sin(t * torch.arange(0, 128, 2) / 128)
time_embedding = torch.cos(t * torch.arange(1, 128, 2) / 128)
time_embedding = torch.cat([time_embedding, time_embedding], dim=0).unsqueeze(0)
# 预测噪声
predicted_noise = model(x_t, time_embedding)
alpha_t = 1 - betas[t.long()]
alpha_t_bar = torch.cumprod(1 - betas, dim=0)[t.long()]
# 计算去噪后的 x_{t-1} (simplified)
x_t_minus_one = (1 / torch.sqrt(alpha_t)) * (x_t - ((1 - alpha_t) / torch.sqrt(1 - alpha_t_bar)) * predicted_noise)
return x_t_minus_one
# 示例：进行一步去噪
# 假设我们已经训练好了一个名为 model 的神经网络
model = SimpleUnet() # 创建一个U-Net模型实例
model.eval() # 设置为评估模式
# 假设 x_t 是一个形状为 (1, 1, 28, 28) 的噪声数据
x_t = torch.randn(1, 1, 28, 28)
t = torch.tensor([500]) # 当前时间步
x_t_minus_one = reverse_diffusion(x_t, t, model, betas)
print("Shape of x_t_minus_one:", x_t_minus_one.shape)解释：
-
model(x_t, t)： 这是调用训练好的神经网络来预测噪声。model接收当前时间步的噪声数据x_t和时间步t作为输入，并输出预测的噪声predicted_noise。 时间步t通常需要进行编码，以便模型能够理解时间信息。这里使用了简单的时间步编码，将时间步归一化后，使用正弦和余弦函数进行编码。实际应用中，可以使用更复杂的编码方式，例如Transformer中的位置编码。 -
alpha_t和alpha_t_bar的计算：alpha_t = 1 - betas[t]计算当前时间步的α_t，即(1 - β_t)。alpha_t_bar = torch.cumprod(1 - betas, dim=0)[t]计算α_t的累积乘积，即(1 - β_1) * (1 - β_2) * ... * (1 - β_t)。 -
x_t_minus_one的计算：x_t_minus_one = (1 / torch.sqrt(alpha_t)) * (x_t - ((1 - alpha_t) / torch.sqrt(1 - alpha_t_bar)) * predicted_noise)使用公式计算去噪后的数据x_{t-1}。 这个公式是基于扩散模型的理论推导得出的，它利用了预测的噪声predicted_noise和α_t以及α_t的累积乘积来估计x_{t-1}。 -
简化的公式: 上述代码中使用的是简化的去噪公式。完整的去噪公式包含方差项的计算，这部分在实际训练中也很重要，但为了简化示例，这里省略了。
7. 完整的采样过程
有了前向扩散和反向扩散的实现，我们就可以进行完整的采样过程了。 采样过程从纯高斯噪声开始，逐步去除噪声，直到得到生成的数据。
def sample(model, image_size, channels, timesteps, betas):
"""
扩散模型采样过程.
Args:
model: 训练好的噪声预测模型.
image_size: 生成图像的大小.
channels: 生成图像的通道数.
timesteps: 总的时间步数.
betas: 噪声时间表.
Returns:
torch.Tensor: 生成的图像.
"""
model.eval()
with torch.no_grad():
# 初始化为纯高斯噪声
x_t = torch.randn((1, channels, image_size, image_size))
# 逐步去噪
for i in reversed(range(timesteps)):
t = torch.full((1,), i, dtype=torch.long)
x_t = reverse_diffusion(x_t, t, model, betas)
return x_t
# 示例：生成一个图像
image_size = 28
channels = 1
timesteps = 1000
generated_image = sample(model, image_size, channels, timesteps, betas)
print("Shape of generated image:", generated_image.shape)8. 训练扩散模型
训练扩散模型的关键是训练噪声预测模型。 训练过程通常使用均方误差（MSE）作为损失函数，目标是使模型预测的噪声尽可能接近真实噪声。
import torch.optim as optim
# 训练参数
epochs = 10
batch_size = 64
learning_rate = 1e-3
# 数据加载 (这里使用随机数据模拟)
train_data = torch.randn(1000, 1, 28, 28)
# 模型和优化器
model = SimpleUnet()
optimizer = optim.Adam(model.parameters(), lr=learning_rate)
mse_loss = nn.MSELoss()
# 训练循环
for epoch in range(epochs):
for i in range(0, len(train_data), batch_size):
# 获取一个batch的数据
x_0 = train_data[i:i + batch_size]
# 随机选择一个时间步
t = torch.randint(0, timesteps, (x_0.shape[0],))
# 前向扩散过程
x_t, noise = forward_diffusion_sample(x_0, t, betas)
# 预测噪声
# 将 t 转换为模型所需的格式 (例如，embedding)
t_float = t.float() / timesteps # 归一化时间步
# 简单的时间步编码, 可以替换为更复杂的编码方式
time_embedding = torch.sin(t_float * torch.arange(0, 128, 2) / 128)
time_embedding = torch.cos(t_float * torch.arange(1, 128, 2) / 128)
time_embedding = torch.cat([time_embedding, time_embedding], dim=0).unsqueeze(0)
predicted_noise = model(x_t, time_embedding)
# 计算损失
loss = mse_loss(predicted_noise, noise)
# 反向传播和优化
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 打印训练信息
if i % 100 == 0:
print(f"Epoch {epoch}, Batch {i}, Loss: {loss.item()}")
print("Training finished!")9. 扩散模型的优势与局限
优势：
- 高质量的生成效果： 扩散模型能够生成非常逼真的图像、音频等数据。
- 训练稳定： 相比GAN等生成模型，扩散模型的训练过程更加稳定。
- 可控性： 扩散模型可以通过调整噪声时间表等参数来控制生成数据的风格。
局限：
- 计算量大： 扩散模型的采样过程需要多次迭代，计算量较大。
- 推理速度慢： 由于采样过程的迭代性，扩散模型的推理速度相对较慢。
10. 总结：代码示例与模型构成
我们讨论了扩散模型的核心思想、数学原理和Python实现，包括噪声时间表的生成、前向扩散过程的采样、反向扩散过程（去噪）的实现以及完整的采样过程。同时，我们还介绍了扩散模型的训练方法以及它的优势与局限。 代码示例展示了如何使用PyTorch实现扩散模型的核心组件，并通过U-Net模型预测噪声。
11. 总结：训练方法与效果影响
训练扩散模型通常使用均方误差（MSE）作为损失函数，通过最小化预测噪声和真实噪声之间的差异来优化模型。扩散模型的生成效果受到多种因素的影响，包括噪声时间表的选择、模型架构的设计、训练数据的质量等。
更多IT精英技术系列讲座，到智猿学院