从零实现DDPM：PyTorch构建扩散模型生成MNIST手写数字

简介基于DDPMDenoising Diffusion Probabilistic Models的PyTorch可运行实现源码面向深度学习研究者和扩散模型初学者覆盖数据预处理、模型构建、训练与采样全流程核心采用U-Net去噪网络并在MNIST数据集上完成训练与图像生成验证。压缩包内共有27个文件以Python源码、模型权重、数据文件及说明文档为主整体大小22.21MB目录结构清晰训练脚本、采样脚本与输出结果分类存放便于直接运行和二次开发。这套源码已有218人学习下载内容涵盖数据集获取、DDPM类设计、训练算法实现、去噪神经网络构建以及实验结果分析并进一步探讨了不同网络架构对生成效果的影响。读者可借助源码与配套说明深入理解扩散模型去噪的数学原理与实现细节掌握完整训练与采样方法通过生成效果与真实图像的对比验证模型有效性为后续优化或项目落地提供可复用的代码基础。 最近把DDPMDenoising Diffusion Probabilistic Models用PyTorch从零实现了一遍跑通了MNIST手写数字生成手上有了一份可以直接复跑的可运行源码。这篇博文就把这份代码连同推导思路一起放出来讲清楚每个模块在干什么、为什么这么写、跑的时候会踩哪些坑。不管你是刚接触扩散模型还是想在本地快速搭一个生成实验照这份代码走一遍基本就能通。我先说结论DDPM的PyTorch实现没有想象中那么复杂核心代码量很小难的是理解几个关键公式怎么落成code。我最初被论文里一长串数学符号劝退过后来耐住性子把公式和代码逐行对照发现真正要写的模块其实就四个噪声调度、降噪网络U-Net、训练循环、采样循环。下面我按这个顺序把代码拆开讲。1. DDPM在做什么先弄清原理再写代码1.1 前向过程给图片逐步加噪DDPM的核心思路很直白训练时我们不断往干净图像里添加高斯噪声直到图像完全变成纯噪声生成时反过来从一个纯噪声出发一步步去掉噪声还原出一张看起来真实的图像。数学上前向加噪可以写成x_t sqrt(alpha_bar_t) * x_0 sqrt(1 - alpha_bar_t) * epsilon这里的x_0是原始图像t是0到T-1之间的步数epsilon是从标准正态分布采样的噪声alpha_bar_t是由一个噪声调度表累积得到的系数。这个公式最大的价值在于不必真的逐帧算1000步加噪只要给定x_0就能一步算出任意时刻t的带噪图像x_t。代码里对应的就是LinearScheduler中的sample_noise函数。可以这样直观理解前向过程相当于往一张照片上不断撒墨水撒到后面墨点覆盖全部原图完全看不清。调度表控制的是“每步撒多少墨”如果每步撒得太少模型会失去学习意义撒得太多图像的原始结构很快就湮灭在噪声里。常用做法是让beta值从1e-4线性增长到0.02对应T取1000步。1.2 反向过程用神经网络学习去噪反向过程才是真正有意思的部分。如果我们能学习到“从带噪图像x_t还原到x_{t-1}”的分布就可以从纯噪声开始迭代生成图像。但直接回归x_{t-1}很难DDPM换了个思路不预测图像而是预测每一步被加进去的噪声。为什么不直接预测图像因为噪声预测任务的target非常干净就是已知的epsilon损失函数用简单的MSE即可稳定训练而直接回归图像在1000步迭代中容易累积误差。论文里证明了只要模型能准确预测噪声就可以推导出x_{t-1}的分布参数然后按公式采样即可。采样公式的核心长这样x_{t-1} 1/sqrt(alpha_t) * (x_t - (1 - alpha_t)/sqrt(1 - alpha_bar_t) * epsilon_theta) sigma_t * z其中epsilon_theta是神经网络预测的噪声z是新的随机噪声t0时不加。生活化理解你手里有一张被墨水污染的字条你训练了一个模型它能根据污染程度判断“哪些是墨点”然后你一点一点减掉墨点最后得到干净字条。2. 代码结构设计为什么这样拆2.1 四个核心模块划分我在设计代码时把整个流程拆成四个独立部分目的是让每一块都能单独调试、替换和扩展。后续如果你想换成DDIM采样、加入Classifier-Free Guidance只需要改对应模块即可不需要重写整套逻辑。模块职责对应实现噪声调度器定义beta、alpha、alpha_bar实现一步加噪LinearScheduler时间嵌入把步数t编码成向量供网络条件输入SinusoidalTimeEmbedding降噪网络接收带噪图像和时间步预测噪声SimpleUNet训练/采样循环负责参数更新和反向生成train()/sample()这里额外解释一下为什么U-Net是扩散模型的主流选择。扩散模型的输入输出都是图像需要保留空间结构U-Net的编码器-解码器结构和跳跃连接可以在下采样过程中不断压缩语义信息、在上采样过程中恢复细节非常适合这类“图像到图像”的预测任务。同时U-Net对输入尺寸的容忍度很高后续换高分辨率数据集不需要改动网络主结构。2.2 运行环境与依赖说明这份代码只依赖torch和torchvision两个包数据使用MNIST首次运行会自动下载。我建议用Python 3.9以上、PyTorch 2.0以上版本CPU也可以训练完整流程只是速度会比较慢如果有NVIDIA GPU训练20个epoch几分钟内就能完成。如果你是conda用户建议新建一个独立环境避免和已有项目共用环境导致版本冲突。3. 核心实现完整可运行的PyTorch源码3.1 噪声调度器与时间嵌入实现先把基础设施写出来。类LinearScheduler做的事情有两件一是预计算所有系数二是提供一个sample_noise方法让训练时直接通过x_0和t得到带噪图像。t在训练中是对每个batch随机采样的因此模型在每步都能学到不同噪声程度的去噪任务这比固定t训练更能提高鲁棒性。import torch import torch.nn as nn import torch.nn.functional as F import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms import math class LinearScheduler: def __init__(self, T1000, beta_start1e-4, beta_end0.02, devicecpu): self.T T self.device device self.betas torch.linspace(beta_start, beta_end, T, devicedevice) self.alphas 1.0 - self.betas self.alpha_bar torch.cumprod(self.alphas, dim0) def sample_noise(self, x_0, t): # t: (B,), 元素取值 [0, T-1] noise torch.randn_like(x_0) sqrt_alpha_bar torch.sqrt(self.alpha_bar[t]).view(-1, 1, 1, 1) sqrt_one_minus_alpha_bar torch.sqrt(1.0 - self.alpha_bar[t]).view(-1, 1, 1, 1) x_t sqrt_alpha_bar * x_0 sqrt_one_minus_alpha_bar * noise return x_t, noise class SinusoidalTimeEmbedding(nn.Module): # 参考Transformer的位置编码把标量t映射成向量 def __init__(self, dim): super().__init__() self.dim dim def forward(self, t): half self.dim // 2 freqs torch.exp(-math.log(10000) * torch.arange(half, devicet.device).float() / half) args t[:, None].float() * freqs[None, :] return torch.cat([torch.sin(args), torch.cos(args)], dim-1)这里有几个容易出错的地方torch.cumprod计算的是累计乘积必须放在alphas上而不是betas上。我最初调试时直接对betas做cumprod导致加噪程度完全错乱。alpha_bar[t]的形状是(B,)要view(-1, 1, 1, 1)才能正确广播到图像张量(B, C, H, W)。时间步t必须用torch.long类型不能是浮点否则索引会报错。3.2 U-Net主干网络实现网络结构我做了简化没有用原论文里ResNetBlockAttention那么复杂的组件但保留了U-Net最核心的骨架三个编码层、两个解码层、跳跃连接。每个卷积块内部都加了GroupNorm和时间步特征注入这已经是扩散模型能稳定训练的最小配置。class ConvBlock(nn.Module): def __init__(self, in_ch, out_ch, emb_ch): super().__init__() self.conv1 nn.Conv2d(in_ch, out_ch, 3, padding1) self.gn1 nn.GroupNorm(8, out_ch) self.conv2 nn.Conv2d(out_ch, out_ch, 3, padding1) self.gn2 nn.GroupNorm(8, out_ch) self.time_mlp nn.Linear(emb_ch, out_ch) self.shortcut nn.Conv2d(in_ch, out_ch, 1) if in_ch ! out_ch else nn.Identity() def forward(self, x, t_emb): h F.silu(self.gn1(self.conv1(x))) h h self.time_mlp(t_emb).unsqueeze(-1).unsqueeze(-1) h F.silu(self.gn2(self.conv2(h))) return self.shortcut(x) h class SimpleUNet(nn.Module): def __init__(self, in_ch1, base_ch64, time_dim256): super().__init__() self.time_embed SinusoidalTimeEmbedding(time_dim) self.time_mlp nn.Sequential( nn.Linear(time_dim, base_ch * 4), nn.SiLU(), nn.Linear(base_ch * 4, base_ch * 4), ) # Encoder self.enc1 ConvBlock(in_ch, base_ch, base_ch * 4) self.down1 nn.Conv2d(base_ch, base_ch * 2, 4, stride2, padding1) self.enc2 ConvBlock(base_ch * 2, base_ch * 2, base_ch * 4) self.down2 nn.Conv2d(base_ch * 2, base_ch * 4, 4, stride2, padding1) self.enc3 ConvBlock(base_ch * 4, base_ch * 4, base_ch * 4) # Decoder self.up1 nn.ConvTranspose2d(base_ch * 4, base_ch * 2, 4, stride2, padding1) self.dec1 ConvBlock(base_ch * 4, base_ch * 2, base_ch * 4) self.up2 nn.ConvTranspose2d(base_ch * 2, base_ch, 4, stride2, padding1) self.dec2 ConvBlock(base_ch * 2, base_ch, base_ch * 4) self.out nn.Conv2d(base_ch, in_ch, 3, padding1) def forward(self, x, t): t_emb self.time_embed(t) t_emb self.time_mlp(t_emb) # Down h1 self.enc1(x, t_emb) h self.down1(h1) h2 self.enc2(h, t_emb) h self.down2(h2) h3 self.enc3(h, t_emb) # Up h self.up1(h3) h torch.cat([h, h2], dim1) h self.dec1(h, t_emb) h self.up2(h) h torch.cat([h, h1], dim1) h self.dec2(h, t_emb) return self.out(h)关于时间嵌入注入方式很多初学者会直接把时间向量拼接到通道维度上但这样会让卷积层输入通道数动态变化非常别扭。更常见的做法是像这里一样把embedding通过一个线性层映射成和特征图通道数相同的向量再以广播方式相加。这相当于在每一层都告诉网络“当前噪声程度是多少”让网络能够按不同力度去噪。3.3 训练循环实现训练循环本身非常简洁核心就是随机取t加噪让U-Net预测噪声和真实噪声算MSE反向传播。原论文用的是简单的MSE实践中也没有必要换成别的损失。每个batch的t都是随机采样的所以模型在每个batch里看到的是不同噪声程度的样本这也是训练稳定的关键。def train(model, scheduler, dataloader, device, epochs20, lr2e-4): model.train() optimizer optim.AdamW(model.parameters(), lrlr) criterion nn.MSELoss() steps_per_epoch len(dataloader) for epoch in range(epochs): total_loss 0.0 for step, (x, _) in enumerate(dataloader): x x.to(device) B x.size(0) t torch.randint(0, scheduler.T, (B,), devicedevice).long() x_t, noise scheduler.sample_noise(x, t) pred_noise model(x_t, t) loss criterion(pred_noise, noise) optimizer.zero_grad() loss.backward() optimizer.step() total_loss loss.item() avg_loss total_loss / steps_per_epoch print(fEpoch {epoch 1:3d}/{epochs} | Loss {avg_loss:.5f})这里有个细节值得注意Optimizer使用AdamW而不是普通Adam。扩散模型的训练中AdamW配合较低学习率2e-4是论文验证过的稳定组合换成SGD很难收敛这个坑我踩过。3.4 采样循环与主入口采样是我们训练模型的最终目的。DDPM的采样是一个从T-1到0的完整循环每一步都调用网络预测噪声再用固定公式更新x。t0时不再添加额外噪声否则最终图像会叠加一层随机扰动导致生成的图像灰蒙蒙、噪点明显。torch.no_grad() def sample(model, scheduler, device, n_samples16, img_size28, channels1): model.eval() x torch.randn(n_samples, channels, img_size, img_size, devicedevice) for t in reversed(range(scheduler.T)): t_tensor torch.full((n_samples,), t, devicedevice, dtypetorch.long) z torch.randn_like(x) if t 0 else torch.zeros_like(x) pred_noise model(x, t_tensor) alpha_t scheduler.alphas[t] alpha_bar_t scheduler.alpha_bar[t] beta_t scheduler.betas[t] coef1 1.0 / torch.sqrt(alpha_t) coef2 (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) x coef1 * (x - coef2 * pred_noise) torch.sqrt(beta_t) * z return x def main(): device torch.device(cuda if torch.cuda.is_available() else cpu) print(Using device:, device) batch_size 128 T 1000 epochs 20 lr 2e-4 transform transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) train_dataset datasets.MNIST(root./data, trainTrue, transformtransform, downloadTrue) dataloader DataLoader(train_dataset, batch_sizebatch_size, shuffleTrue, num_workers2, drop_lastTrue) scheduler LinearScheduler(TT, devicedevice) model SimpleUNet(in_ch1, base_ch64, time_dim256).to(device) print(Model params:, sum(p.numel() for p in model.parameters())) train(model, scheduler, dataloader, device, epochsepochs, lrlr) torch.save(model.state_dict(), ddpm_mnist.pth) samples sample(model, scheduler, device, n_samples64, img_size28, channels1) samples torch.clamp((samples 1) / 2, 0, 1) from torchvision.utils import save_image save_image(samples, ddpm_samples.png, nrow8) print(Saved samples to ddpm_samples.png) if __name__ __main__: main()主入口里我把数据归一化到[-1, 1]这一点非常关键。MNIST默认的ToTensor把像素缩放到[0, 1]但扩散模型训练时通常需要输入在[-1, 1]范围这样才能和标准正态分布噪声在数值上匹配。保存图片时再映射回[0, 1]并clamp即可。注意torch.randn采样的是标准正态分布如果训练数据是[0,1]范围模型就需要输出负值来抵消噪声但网络表达能力有限生成的图像会整体偏亮且细节差。归一化到[-1,1]是扩散模型的基本操盘手法。4. 训练效果与超参解读4.1 MNIST上的收敛表现我在自己的机器上单张NVIDIA GPU用这份代码训练了20个epochbatch size128T1000实测loss从最初的约0.35下降到0.02以下采样出来的手写数字已经清晰可辨。第5个epoch左右就能看到模糊的数字轮廓到第15个epoch后大部分采样结果接近真实MNIST风格。如果你在CPU上跑20个epoch可能要几个小时到十几个小时不等。作为快速验证可以把T降到300、epoch降到5先确认代码链路通了再跑完整训练。T越小训练越快但采样质量会下降这个取舍需要根据实验目的来定。4.2 关键超参会影响什么我把在调试过程中实际影响比较大的几个超参整理成表方便你对照调整超参默认值影响说明T1000越大生成细节越好但训练和采样耗时线性增加beta_start / beta_end1e-4 / 0.02范围过窄会让前向噪声不足过宽会让信号过早湮灭lr2e-4调大容易震荡调小收敛慢AdamW下这个值很稳batch_size128影响梯度稳定性和显存占用显存不足时优先降到64或32base_ch64决定网络宽度想提高表现可以翻倍到128但参数量约膨胀4倍epochs20MNIST上用20个epoch已经够用复杂数据集需要更多训练前建议打印一下模型参数量这份简化版U-Net约2300万参数比原论文的缩小了很多但生成MNIST已经绰绰有余。如果你跑CIFAR-10之类的复杂数据集需要在网络深度、通道数和Attention机制上做增强。5. 常见问题与排查技巧实录5.1 训练Loss不下降怎么办先看Loss量级。如果第一个epoch就在0.3左右并且缓慢下降属于正常如果Loss完全不动或者直接变成NaN优先排查三件事一是输入数据是否在[-1,1]二是t是否传成浮点类型导致索引出问题三是学习率是否过高超过1e-3很容易震荡。另外别忘了检查网络输出形状是否和噪声一致我曾经把U-Net的out通道设成3在MNIST上跑了一整晚才发现白白浪费了时间。5.2 采样结果全是噪声这个是最常见的问题。如果你是训练没几个epoch就急着采样那大概率只是欠拟合多训几个epoch再看。如果训练Loss已经很低但采样还是纯噪声重点检查采样循环里的公式系数是不是误用了alpha_bar[t]而不是alpha[t]计算coef1coef2的分母是否正确取了sqrt(1 - alpha_bar_t)t0时是否加了额外噪声模型是否忘记调用load_state_dict或者加载了错误的权重文件。一个实用的调试技巧把T临时调成50或100采样速度会快很多可以快速验证逻辑是否正确如果小T下能生成大致轮廓说明采样公式没错只是训练精度不够。5.3 显存不足或者训练速度太慢降低batch size永远是最快的解法。把128改成64、32Loss曲线会稍微抖一点但不影响整体收敛。其次是降低base_ch把64改成32参数量会大幅下降训练速度明显提升。还有一招是开启混合精度训练PyTorch自带torch.cuda.amp对显存占用和速度都有帮助。如果你只能在CPU上跑建议缩小网络把base_ch设为32T降到300epoch控制在5以内。这套配置虽然不是完整训练但足以跑通流程、理解代码。5.4 关于环境配置的几句话很多人在环境搭配上花了不少时间我个人的习惯是先装好Anaconda然后创建独立虚拟环境再根据机器是否有NVIDIA GPU选择对应的PyTorch安装命令。CPU版本安装起来最简单直接装默认版本即可GPU版本需要注意CUDA驱动和PyTorch的对应关系装错版本最常见的报错就是“CUDA unavailable”或者“driver version insufficient”。如果是在Ubuntu服务器上训练建议先输入nvidia-smi确认驱动支持的CUDA版本再进入PyTorch官网选择匹配的安装命令。Windows用户要注意DataLoader里的num_workers设置成大于0可能会遇到多进程报错直接改为0即可。我在实际跑这个项目的过程中最大的体会是扩散模型论文看起来门槛高但代码实现的门槛其实远低于预期。关键是把“前向加噪、噪声预测、反向采样”这三段逻辑彻底理解清楚再配合一份可运行的代码对照调试很快就能建立直观认识。最后再分享一个小技巧训练过程中每隔几个epoch就手动采样保存一张图看着图像从纯噪声一步步变成清晰数字比盯着Loss数值更能判断模型的真实状态。这份代码你也可以接着往上扩展DDIM采样、Classifier-Free Guidance甚至换成CIFAR-10数据集做条件生成希望你也能跑出自己的第一张扩散生成图。本文还有配套的精品资源点击获取