从零实现Diffusion模型：基于PyTorch的DDPM实战指南

原创

于 2026-03-07 00:45:28 发布 · 441 阅读

标签

1. 从零理解DDPM：它到底是什么，为什么能“无中生有”？

你可能已经见过很多由AI生成的、以假乱真的图片了，从写实的人像到奇幻的场景，这些图片很多都出自一种叫做“扩散模型”的技术。而我们要动手实现的DDPM，正是扩散模型家族里最经典、也最适合入门的一位成员。我第一次接触DDPM时，感觉它的思想特别巧妙，它不像GAN那样让两个网络互相“打架”，也不像VAE那样去学习一个复杂的分布。它的核心灵感，其实来源于一个非常自然的物理现象：扩散。

想象一滴墨水掉进一杯清水里。一开始，墨水的浓度很高，聚集在一处。随着时间的推移，墨水分子会逐渐从高浓度区域向低浓度区域运动，最终均匀地散布在整个水杯中，清水变成了一杯颜色均匀的淡墨水。这个过程就是“扩散”。DDPM的“正向过程”就是在模拟这个“滴墨入水”的过程：它把一张清晰的图片（墨水团）通过一步步添加噪声，最终变成一张完全随机的、如同电视雪花屏一样的纯噪声图片（均匀的淡墨水）。

那么，AI怎么从噪声中生成图片呢？这就是“逆向过程”的神奇之处了。模型要学习的，就是如何“倒放”这个扩散过程。如果我们能知道每一步噪声扩散的“反方向”，就能从一杯均匀的淡墨水中，让墨水分子重新聚集起来，变回最初的那滴墨水。当然，在数学和计算上，直接精确地逆转这个过程是极其困难的。DDPM的聪明之处在于，它不直接预测去噪后的清晰图片，而是预测当前噪声图片中所包含的“噪声”是什么。这就像是你知道一杯淡墨水是由一滴浓墨化开的，虽然无法直接变回那滴墨，但你可以估算出“化开”这个动作本身。通过不断减去预测出的噪声，图片就逐渐从混沌中浮现出来。

这个思路让训练变得非常稳定。我刚开始尝试复现时，最头疼的就是GAN那种难以捉摸的训练崩溃问题。而DDPM的损失函数简单直接：就是让模型预测的噪声，和我们在正向过程中实际加入的噪声，尽可能一样。模型的目标变得纯粹而单一，训练过程也就稳当多了。接下来，我们就用PyTorch，一步步把这个有趣的想法变成可以运行的代码。

2. 实战第一步：搭建开发环境与准备数据

动手之前，得先把“厨房”收拾好。我强烈建议使用Anaconda来管理Python环境，它能帮你轻松处理各种依赖包版本冲突的问题，这可是我踩过不少坑才得出的经验。

打开你的终端或Anaconda Prompt，执行下面的命令来创建一个专门用于本项目的环境：

conda create -n ddpm_pytorch python=3.9
conda activate ddpm_pytorch

接着，安装核心的PyTorch。这里要注意，你需要根据自己电脑是否有NVIDIA显卡，以及CUDA的版本，去PyTorch官网选择对应的安装命令。有显卡的话训练速度会快很多。例如，对于CUDA 11.8，可以这样安装：

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

如果没有显卡，就安装CPU版本：

pip install torch torchvision torchaudio

然后，我们安装一些必要的辅助库，比如用于画图的matplotlib和用于更便捷数据操作的numpy：

pip install matplotlib numpy

环境配好了，接下来是数据。为了让实验快速跑起来，我们选择一个规模适中、又比较有趣的数据集：斯坦福汽车数据集（StanfordCars）。这个数据集包含了196类不同品牌的汽车图片，足够让我们的模型学习到“车”这个概念的基本结构。使用torchvision可以非常方便地下载和加载它。

import torch
import torchvision
import matplotlib.pyplot as plt
from torchvision import transforms
from torch.utils.data import DataLoader
import numpy as np

# 设置设备
device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

# 定义图片大小和批次大小
IMG_SIZE = 64  # 为了训练效率，我们将图片统一缩放到64x64
BATCH_SIZE = 128

def load_transformed_dataset():
    """
    定义数据预处理流程并加载数据集。
    预处理包括：调整大小、随机水平翻转（数据增强）、转换为张量、并将像素值从[0,1]映射到[-1,1]。
    """
    data_transforms = [
        transforms.Resize((IMG_SIZE, IMG_SIZE)),
        transforms.RandomHorizontalFlip(), # 简单增强，让模型更鲁棒
        transforms.ToTensor(), # 将PIL图像转换为[0,1]范围的张量
        transforms.Lambda(lambda t: (t * 2) - 1) # 将[0,1]映射到[-1,1]，这对模型训练更稳定
    ]
    data_transform = transforms.Compose(data_transforms)

    # 下载并加载训练集和测试集
    train_dataset = torchvision.datasets.StanfordCars(root="./data", download=True,
                                                       transform=data_transform)
    test_dataset = torchvision.datasets.StanfordCars(root="./data", download=True,
                                                      transform=data_transform, split='test')
    # 合并起来使用，对于简单的生成任务数据越多越好
    return torch.utils.data.ConcatDataset([train_dataset, test_dataset])

# 加载数据
dataset = load_transformed_dataset()
dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=True, drop_last=True, num_workers=4)

# 写一个辅助函数来看看我们的数据长什么样
def show_tensor_image(image_tensor):
    """
    将[-1,1]范围的张量图像转换回PIL图像并显示。
    """
    reverse_transforms = transforms.Compose([
        transforms.Lambda(lambda t: (t + 1) / 2), # 从[-1,1]映射回[0,1]
        transforms.Lambda(lambda t: t.permute(1, 2, 0)), # 从(C,H,W)转换为(H,W,C)
        transforms.Lambda(lambda t: t * 255.), # 乘255
        transforms.Lambda(lambda t: t.cpu().numpy().astype(np.uint8)), # 转numpy数组
    ])
    # 如果输入是批次数据，只取第一张
    if len(image_tensor.shape) == 4:
        image_tensor = image_tensor[0]
    image_np = reverse_transforms(image_tensor)
    plt.imshow(image_np)
    plt.axis('off')

# 可视

最低0.47元/天解锁文章