Diffusion

原创已于 2025-05-06 07:06:03 修改 · 689 阅读

13 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #python #人工智能

于 2025-03-01 15:24:29 首次发布

自动驾驶---人工智能专栏收录该内容

2 篇文章

订阅专栏

Vanilla Difussion

UNet 介绍

扩散过程的核心思想
扩散模型的核心思想是通过逐步添加噪声将数据（如图像）从原始状态（x_0）转变为纯噪声（x_T），然后再通过反向过程从噪声中恢复出原始数据。这个过程分为两个阶段：

前向过程（Forward Process）：逐步添加噪声。

反向过程（Reverse Process）：逐步去除噪声。

timesteps 定义了前向过程和反向过程中有多少个时间步。

总时间步数 timesteps 的作用
定义噪声添加的粒度：

timesteps 越大，噪声添加的步骤越多，每一步添加的噪声量越小。

timesteps 越小，噪声添加的步骤越少，每一步添加的噪声量越大。

控制扩散过程的平滑性：

较大的 timesteps 使得扩散过程更加平滑，适合生成高质量的数据。

较小的 timesteps 可能导致扩散过程不够平滑，生成的数据质量较差。

前向过程中的 timesteps
在前向过程中，数据从原始状态 x_0 逐步变为噪声 x_T，每一步的噪声量由噪声调度（如 self.betas）控制。timesteps 决定了：

每一步的噪声强度：self.betas 是一个长度为 timesteps 的张量，表示每一步的噪声强度。

每一步的状态：x_t 是第 t 步的数据状态，t 的取值范围是 [0, timesteps-1]。

示例
假设 timesteps=1000：

前向过程：

从 x_0 到 x_{1000}，每一步添加少量噪声。

噪声强度由 self.betas 控制，self.betas 是一个长度为 1000 的张量。

反向过程：

从 x_{1000} 开始，逐步去除噪声，恢复出 x_0。

每一步去除的噪声由模型预测。

使用 Diffusion 生成图像的步骤为，这里以 MNIST 为例：
首先要训练网络，根据噪声生成图像的能力：

S1：前向传播为 x_0 添加噪声
这个郭过程不关乎任何网络训练，只是根据 timestamp 给图像添加噪声

def forward_process(self, x_0, t):
        """前向过程：向输入数据 x_0 添加噪声"""
        noise = torch.randn_like(x_0)  # 生成与 x_0 形状相同的随机噪声
        alpha_bar_t = self.alpha_bars[t].view(-1, 1, 1, 1)  # 获取时间步 t 对应的 alpha_bar，并调整形状以匹配 x_0
        x_t = torch.sqrt(alpha_bar_t) * x_0 + torch.sqrt(1 - alpha_bar_t) * noise  # 添加噪声
        return x_t, noise

S2：预测噪声图像的噪声分布
这里使用 UNet，传入刚刚添加了噪声的图像，预测原图像的噪声分布

def reverse_process(self, x_t, t):
        """反向过程：使用模型预测噪声"""
        predicted_noise = self.model(x_t)  # 通过 U-Net 预测噪声
        return predicted_noise

S3：loss 回传
根据刚才的步骤，需要将生成噪声与预测噪声两个步骤联合起来，训练 UNet 网络，这里损失函数传入生成的噪声真值和 UNet 预测的噪声

def train_step(self, x_0, optimizer):
       """训练步骤：训练模型预测噪声"""
       self.model.train()  # 设置模型为训练模式
       optimizer.zero_grad()  # 清空梯度
       t = torch.randint(0, self.timesteps, (x_0.shape[0],))  # 随机采样时间步
       x_t, noise = self.forward_process(x_0, t)  # 前向过程：添加噪声
       predicted_noise = self.reverse_process(x_t, t)  # 反向过程：预测噪声
       loss = F.mse_loss(predicted_noise, noise)  # 计算预测噪声与真实噪声的均方误差
       loss.backward()  # 反向传播
       optimizer.step()  # 更新模型参数
       return loss.item()  # 返回损失值

S4：传入噪声，生成样本
当网络训练完成后传入噪声，生成 MNIST 数据

def sample(self, num_samples=1, image_size=(1, 28, 28)): # 生成图像
        """采样生成图像：从随机噪声开始，逐步去噪生成图像"""
        self.model.eval()  # 设置模型为评估模式
        with torch.no_grad():  # 禁用梯度计算
            x_t = torch.randn((num_samples, *image_size))  # 先生成一张噪声图像
            for t in reversed(range(self.timesteps)):  # 从最后一个时间步逐步去噪
                t_tensor = torch.full((num_samples,), t, dtype=torch.long)  # 创建时间步张量
                predicted_noise = self.reverse_process(x_t, t_tensor)  # 预测噪声
                alpha_t = self.alphas[t]  # 获取当前时间步的 alpha
                alpha_bar_t = self.alpha_bars[t]  # 获取当前时间步的 alpha_bar
                beta_t = self.betas[t]  # 获取当前时间步的 beta
                if t > 0:
                    noise = torch.randn_like(x_t)  # 如果不是最后一步，添加随机噪声
                else:
                    noise = torch.zeros_like(x_t)  # 如果是最后一步，不添加噪声
                # 更新 x_t：根据预测噪声和噪声调度公式逐步去噪
                x_t = (1 / torch.sqrt(alpha_t)) * (x_t - ((1 - alpha_t) / torch.sqrt(1 - alpha_bar_t)) * predicted_noise) + torch.sqrt(beta_t) * noise
        return x_t  # 返回生成的图像

完整代码

import torch
import torch.nn as nn
import torch.nn.functional as F
import matplotlib.pyplot as plt
from torchvision import transforms
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST

# 定义 U-Net 网络（简化版）
class UNet(nn.Module):
    def __init__(self, in_channels=1, out_channels=1):
        super(UNet, self).__init__()
        # 编码器部分：逐步下采样，提取特征
        self.encoder = nn.Sequential(
            nn.Conv2d(in_channels, 64, kernel_size=3, stride=1, padding=1),  # 卷积层，输入通道到64通道
            nn.ReLU(),  # 激活函数
            nn.Conv2d(64, 128, kernel_size=3, stride=2, padding=1),  # 卷积层，64通道到128通道，步长为2（下采样）
            nn.ReLU(),  # 激活函数
        )
        # 解码器部分：逐步上采样，重建图像
        self.decoder = nn.Sequential(
            nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),  # 转置卷积层，128通道到64通道，步长为2（上采样）
            nn.ReLU(),  # 激活函数
            nn.Conv2d(64, out_channels, kernel_size=3, stride=1, padding=1),  # 卷积层，64通道到输出通道
        )

    def forward(self, x):
        # 前向传播：先通过编码器，再通过解码器
        x = self.encoder(x)
        x = self.decoder(x)
        return x

# 定义扩散模型
class VanillaDiffusion:
    def __init__(self, model, timesteps=1000, beta_start=1e-4, beta_end=0.02):
        self.model = model  # U-Net 模型
        self.timesteps = timesteps  # 扩散过程的总时间步数
        self.betas = torch.linspace(beta_start, beta_end, timesteps)  # 噪声调度：线性增加的噪声强度
        self.alphas = 1.0 - self.betas  # alpha = 1 - beta
        self.alpha_bars = torch.cumprod(self.alphas, dim=0)  # alpha_bar 是 alpha 的累积乘积，用于计算前向过程的噪声

    def forward_process(self, x_0, t):
        """前向过程：向输入数据 x_0 添加噪声"""
        noise = torch.randn_like(x_0)  # 生成与 x_0 形状相同的随机噪声
        alpha_bar_t = self.alpha_bars[t].view(-1, 1, 1, 1)  # 获取时间步 t 对应的 alpha_bar，并调整形状以匹配 x_0
        x_t = torch.sqrt(alpha_bar_t) * x_0 + torch.sqrt(1 - alpha_bar_t) * noise  # 添加噪声
        return x_t, noise

    def reverse_process(self, x_t, t):
        """反向过程：使用模型预测噪声"""
        predicted_noise = self.model(x_t)  # 通过 U-Net 预测噪声
        return predicted_noise

    def train_step(self, x_0, optimizer):
        """训练步骤：训练模型预测噪声"""
        self.model.train()  # 设置模型为训练模式
        optimizer.zero_grad()  # 清空梯度
        t = torch.randint(0, self.timesteps, (x_0.shape[0],))  # 随机采样时间步
        x_t, noise = self.forward_process(x_0, t)  # 前向过程：添加噪声
        predicted_noise = self.reverse_process(x_t, t)  # 反向过程：预测噪声
        loss = F.mse_loss(predicted_noise, noise)  # 计算预测噪声与真实噪声的均方误差
        loss.backward()  # 反向传播
        optimizer.step()  # 更新模型参数
        return loss.item()  # 返回损失值

    def sample(self, num_samples=1, image_size=(1, 28, 28)): # 生成图像
        """采样生成图像：从随机噪声开始，逐步去噪生成图像"""
        self.model.eval()  # 设置模型为评估模式
        with torch.no_grad():  # 禁用梯度计算
            x_t = torch.randn((num_samples, *image_size))  # 先生成一张噪声图像
            for t in reversed(range(self.timesteps)):  # 从最后一个时间步逐步去噪
                t_tensor = torch.full((num_samples,), t, dtype=torch.long)  # 创建时间步张量
                predicted_noise = self.reverse_process(x_t, t_tensor)  # 预测噪声
                alpha_t = self.alphas[t]  # 获取当前时间步的 alpha
                alpha_bar_t = self.alpha_bars[t]  # 获取当前时间步的 alpha_bar
                beta_t = self.betas[t]  # 获取当前时间步的 beta
                if t > 0:
                    noise = torch.randn_like(x_t)  # 如果不是最后一步，添加随机噪声
                else:
                    noise = torch.zeros_like(x_t)  # 如果是最后一步，不添加噪声
                # 更新 x_t：根据预测噪声和噪声调度公式逐步去噪
                x_t = (1 / torch.sqrt(alpha_t)) * (x_t - ((1 - alpha_t) / torch.sqrt(1 - alpha_bar_t)) * predicted_noise) + torch.sqrt(beta_t) * noise
        return x_t  # 返回生成的图像

# 数据加载
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,))])  # 数据预处理：将图像转换为张量并归一化
dataset = MNIST(root="./data", train=True, download=True, transform=transform)  # 加载 MNIST 数据集
dataloader = DataLoader(dataset, batch_size=128, shuffle=True)  # 创建数据加载器

# 初始化模型和优化器
model = UNet()  # 创建 U-Net 模型
diffusion = VanillaDiffusion(model)  # 创建扩散模型
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)  # 使用 Adam 优化器

# 训练
epochs = 10  # 训练轮数
for epoch in range(epochs):
    for batch in dataloader:
        x_0, _ = batch  # 获取输入数据
        loss = diffusion.train_step(x_0, optimizer)  # 训练步骤
    print(f"Epoch {epoch + 1}, Loss: {loss}")  # 打印每轮的损失值

# 生成图像
generated_images = diffusion.sample(num_samples=16)  # 生成 16 张图像
generated_images = (generated_images + 1) / 2  # 反归一化：将图像从 [-1, 1] 范围转换到 [0, 1] 范围

# 保存生成的图像到本地
utils.save_image(generated_images, "generated_images.png", nrow=4)  # 将图像保存为 PNG 文件，每行显示 4 张图像
print("Generated images saved to 'generated_images.png'")  # 打印保存成功的消息