手写Python扩散模型：从零训练S型曲线生成器

简介：这是一份面向初学者的扩散模型入门演示资源，以生成S型曲线为实例，帮助理解扩散模型从随机噪声逐步还原数据分布的核心思想。资源适合刚接触生成模型、希望将理论与代码对照学习的读者，也适合作为课堂或自学项目的基础样例。压缩包共8个文件，约9.74MB，其中ipynb为可运行的Notebook代码，gif直观展示扩散与生成过程，xml和iml为项目配置文件，gitignore用于工程管理，整体结构清晰、便于直接打开运行。目前已有1600余人浏览学习。通过运行该demo，可以观察S型曲线在扩散与去噪过程中的变化，尝试调整时间步长或噪声参数来对比生成效果，从而加深对前向加噪、反向采样等概念的理解。对于希望快速建立扩散模型直觉并动手实践的学习者来说，这是一份轻量但完整的入门素材。 扩散模型（diffusion model）这个名词这两年几乎成了生成模型的代名词，但很多初学者一上来就被一堆数学公式和UNet结构劝退。我前几天给一个师弟写了个小demo，只有几百行Python代码，在CPU上跑两分钟，就能训练出一个完整的扩散模型，用来生成一条S型曲线，也就是sigmoid函数形状的分布。跑完之后他跟我说：“原来扩散模型就这么回事，之前看论文完全没概念。”

这个demo说通俗点就是：让模型从一堆随机噪声点里，自己慢慢“长”出一条S型曲线。整个过程把扩散模型最核心的加噪、去噪、采样三步完整走了一遍，非常适合想真正理解DDPM而不是只调包的朋友。不需要GPU，不需要下载数据集，公式只需要两个核心的，剩下的全是直观的代码和图像。这篇文章我就把这个demo的完整思路、原理、代码和踩过的坑全部写出来，照着敲一遍，你会比看十篇科普更懂扩散模型。

1. 为什么我推荐用S型曲线入门扩散模型

1.1 先建立直觉再啃公式

扩散模型本质上就是一个“先破坏再修复”的游戏。训练阶段先往干净的数据上逐步添加高斯噪声，直到数据完全变成一团纯噪声，这叫前向过程；然后训练一个神经网络，让它学会从任意程度的噪声里把原始数据还原出来，这叫反向过程。推理阶段就更有意思了：从一个纯噪声开始，让网络一步步去噪，最后还原出一个全新的样本。

用生活化的类比来解释：想象你有一杯牛奶，每次往里加一点水，搅匀后再加一点水，重复几百次后牛奶变成了完全透明的清水。扩散模型做的事情就是训练一个AI，让它学会“逆着这个过程”，从一杯清水一步步把牛奶还原出来。实际训练中AI不认识牛奶也没关系，它只需要学会在给定当前混合液体的状态下，估算出这一步刚加进去的那点水是长什么样的——这就是预测噪声。

S型曲线demo做的事情完全不同，但机制完全一样：目标是让模型学会生成一个二维点集，这些点分布在sigmoid函数形状的带状区域上。你可以把这个带状区域当成“牛奶”，把高斯噪声当成“水”。训练完成后，从一个随机噪声点云开始，模型能逐步把它去噪成漂亮的S型点云。整个过程完全可可视化，每一步都看得见，这是图像demo给不了的直观感。

1.2 曲线数据比图像数据友好太多

图像数据动辄几十万维，网络结构必须上UNet，训练必须用GPU，光是把环境跑通就让很多初学者倒在起跑线前。但这个demo里每个样本只是二维坐标点，网络就是一个几百K参数的MLP（多层感知机），CPU上几十秒到两分钟就能完成训练。

我做了个对比表格，差别一目了然：

这个demo虽然看起来小，但它把扩散模型的核心骨架完整保留了：噪声调度器、加噪过程、时间步嵌入、噪声预测网络、反向采样、结果评估，一个都不少。所以它完全可以直接作为后续学习图像扩散模型的“最小可运行骨架”，你理解了这版代码里的每个组件，再去看Stable Diffusion、DDIM、LDM之类的进阶内容，就不会觉得它们是黑盒了。

1.3 这个demo覆盖了DDPM的完整链路

我特意保留了扩散模型的所有关键环节，没有做任何“省事”的简化——训练目标是最经典的噪声预测（epsilon预测），采样过程是完整的DDPM逐步去噪，时间步用了正弦位置编码。只是把“高维图像”降维成了“二维点云”，把“UNet”换成了“MLP”。

这意味着你在这个demo里学到的每一个概念，迁移到正式项目里都是成立的。比如模型输入必须是加噪后的数据加上对应的时间步，比如采样时刻末步不能加随机噪声，这些都是真实扩散模型里的核心细节，提前在小demo里踩坑，以后就不会踩第二次。

2. 扩散模型核心机制深度拆解

2.1 前向过程：怎么把一条S型曲线搅成噪声

扩散模型的第一步是前向过程，给定原始样本x0，按一个预设的噪声调度器逐步添加高斯噪声，直到样本完全变成标准正态分布。

设总步数为T（代码里我用200，原版DDPM是1000），每一步的噪声强度由beta_t控制。从x0直接获取任意时刻t的加噪结果，不需要一步一步递推，有一句闭式公式：

x_t = sqrt(alpha_bar_t) * x0 + sqrt(1 - alpha_bar_t) * epsilon

其中alpha_bar_t是前t步的alpha累乘，alpha_t = 1 - beta_t，epsilon是标准正态噪声。这个公式的物理含义很直白：第t步的加噪样本，就是原始样本和纯噪声按一定比例混合的结果。t越靠近T，混合中噪声占比越高，样本被污染得越严重。

代码实现如下：

import torch def q_sample(x0, t, alpha_bar, noise): # x0: [batch, 2] 原始二维点 # t: [batch] 时间步 # noise: [batch, 2] 标准正态噪声 sqrt_alpha_bar = torch.sqrt(alpha_bar[t])[:, None] sqrt_one_minus_alpha_bar = torch.sqrt(1 - alpha_bar[t])[:, None] return sqrt_alpha_bar * x0 + sqrt_one_minus_alpha_bar * noise

2.2 反向过程：让神经网络学会“去污”

前向过程把数据搅成噪声之后，我们需要训练一个神经网络来学习它的逆过程。网络输入是加噪样本x_t和时间步t，输出是网络预测的噪声epsilon_theta(x_t, t)。

训练目标出奇简洁，就是让预测噪声逼近真实噪声，用均方误差度量：

L = E[|| epsilon - epsilon_theta(x_t, t) ||^2]

再次用类比解释：加噪过程好比是你往一道菜里加盐，你知道加了“多少克盐”，但网络不知道。你现在把加了盐的菜和“加盐次数”告诉网络，让它猜这轮实际加了多少克盐。只要它猜得准，将来从纯盐（纯噪声）开始还原一道菜时，就能一步步“剥离”出正确的盐量，也就是还原出原始数据。

这里有个很关键的设计：为什么是预测噪声，而不是直接预测原始数据x0？直接预测x0看起来更直接，但实验证明预测噪声更稳定。原因在于，加噪公式里x_t和x0、epsilon是线性混合关系，网络直接学习噪声分量在数学上等价于学习数据的得分函数（score function），这是扩散模型理论的基石之一。作为初学者不需要完整啃推导，但记住“预测噪声”这个约定俗成的范式，对你阅读后续论文很有帮助。

2.3 采样：从纯噪声开始“雕刻”数据

训练完成后，生成过程就是从纯噪声开始逐步去噪。设当前时刻为t，上一时刻为t-1，采样公式是：

x_{t-1} = 1 / sqrt(alpha_t) * ( x_t - (1 - alpha_t) / sqrt(1 - alpha_bar_t) * epsilon_theta(x_t, t) ) + sigma_t * z

这个公式看着吓人，实际含义分两部分理解：前半部分是“去噪均值”，根据网络预测的噪声，把x_t朝原始数据的方向拉一拉；后半部分是“重新注入噪声”，z是标准正态噪声，sigma_t是当前步的噪声强度系数。重新注入噪声是为了保证生成分布的正确性，防止一下子把所有噪声信息抹掉导致数据坍缩成确定值——就好比做菜加盐要一点点来，一次到位容易淡咸失手。

注意当t=1到达最后一步时，z必须置为0，也就是最后一步只做去噪，不再注入新噪声，这样生成出的点才干净。

2.4 为什么要加时间嵌入

网络的核心输入除了x_t本身，还有一个极其重要、初学者特别容易忽略的输入——时间步t。

不同t时刻，x_t的噪声程度完全不同。t靠近0时数据基本是干净的，t靠近T时基本是纯噪声。如果网络不知道t，它就无法区分当前数据处于“轻污染”还是“重污染”状态，只能对所有输入拟合同一个映射，结果就是什么都学不好，训练loss高居不下，生成出来的全是模糊的一团。

解决办法是把时间步t编码成向量，再接进网络。我沿用DDPM里的正弦位置编码方式，把t转换成一个64维的向量，再用一个小MLP映射到隐藏层维度，然后加到特征里。代码实现如下：

import torch import torch.nn as nn import math class SinusoidalTimeEmbedding(nn.Module): def __init__(self, dim): super().__init__() self.dim = dim def forward(self, t): # t: [batch] half_dim = self.dim // 2 emb = math.log(10000) / (half_dim - 1) emb = torch.exp(torch.arange(half_dim, dtype=torch.float32) * -emb) emb = t[:, None] * emb[None, :] return torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)

没有时间嵌入的扩散模型，相当于让一个盲人去修复一幅画，他连画被糟蹋到了什么程度都不知道，自然无从下手。

3. 完整实现：从数据准备到采样

3.1 环境与数据准备

这个demo只依赖三个库：PyTorch用于搭建和训练网络，NumPy用于数据生成和数学运算，Matplotlib用于结果可视化。安装命令：

pip install torch numpy matplotlib

数据准备上我用了很取巧的方式：不预先构建固定数据集，而是在训练时实时生成样本。方法是从[-6, 6]区间均匀随机采样x坐标，带入sigmoid函数计算对应的y值，再加一点高斯噪声。这样数据是无限的，网络不会过拟合，而且每个batch都能覆盖完整的曲线范围。

import numpy as np def sigmoid(x): return 1 / (1 + np.exp(-x)) def generate_batch(batch_size=1024, noise_scale=0.05): x = np.random.uniform(-6, 6, size=batch_size) y_true = sigmoid(x) y = y_true + np.random.randn(batch_size) * noise_scale return torch.tensor(np.stack([x, y], axis=1), dtype=torch.float32)

为什么用sigmoid曲线的点云而不是一条光滑曲线？因为扩散模型本质上是学习概率分布的，让模型学习“S型带状分布”比“S型细线”更容易，也更贴近真实世界的生成任务——真实数据通常都带有噪声和多样性。噪声缩放系数定在0.05，既能体现数据的分布宽度，又不会把S形状淹没了。

3.2 噪声调度器与加噪

DDPM的噪声调度器用的是线性衰减：beta从0.0001均匀增加到0.02。我用T=200步，比原版1000步少了五倍，好处是训练和采样都快很多，坏处是单步去噪量更大，模型压力略高，但对这个二维小任务完全够用。

def create_noise_schedule(T=200, beta_min=0.0001, beta_max=0.02): betas = torch.linspace(beta_min, beta_max, T) alphas = 1.0 - betas alpha_bar = torch.cumprod(alphas, dim=0) return betas, alphas, alpha_bar

有个细节值得提一下：beta的范围不能随便调大。如果beta_max超过0.1，意味着最后几步的数据几乎瞬间被噪声完全覆盖，模型学习难度会骤增。反之如果beta_max太小，噪声不够强，生成结果的多样性会明显下降。0.0001到0.02这个区间是经过大量实践验证的经典取值，入门阶段直接用就好。

3.3 去噪网络：MLP加时间嵌入

因为输入是二维坐标，不需要复杂的UNet结构，一个四层MLP足够拟合。网络的核心结构是：把时间嵌入向量作为条件，拼接到每一层特征中。我用了最简单的方式——每层都加时间嵌入向量，而不是只在输入层拼一次。这样做的好处是时间信息能更深入地参与各层特征变换，让网络在每一个抽象层级都能感知当前噪声程度。

import torch.nn as nn class Denoiser(nn.Module): def __init__(self, hidden_dim=128, time_dim=64): super().__init__() self.time_emb = SinusoidalTimeEmbedding(time_dim) self.time_proj = nn.Sequential( nn.Linear(time_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim), ) self.net = nn.Sequential( nn.Linear(2 + hidden_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, hidden_dim), nn.SiLU(), nn.Linear(hidden_dim, 2), ) def forward(self, x, t): # x: [batch, 2] 加噪后的坐标 # t: [batch] 时间步 t_emb = self.time_proj(self.time_emb(t)) h = torch.cat([x, t_emb], dim=-1) return self.net(h)

激活函数我选的是SiLU，也就是swish，它在深度网络里的表现通常比ReLU更平滑，对于二维连续输出任务更友好。你也可以换成ReLU试一下，会发现损失曲线波动略大，这就是激活函数选择对训练稳定性的直观影响。

3.4 训练主循环

训练算法非常简单，五步完成：生成一批数据，随机采样时间步，按时间步对数据进行加噪，让网络预测噪声，计算MSE并反向传播。下面是完整训练代码：

def train(model, optimizer, steps=6000, batch_size=1024, T=200, alpha_bar=None): model.train() losses = [] for step in range(steps): x0 = generate_batch(batch_size) t = torch.randint(1, T, (batch_size,), dtype=torch.float32) noise = torch.randn_like(x0) xt = q_sample(x0, t, alpha_bar, noise) pred_noise = model(xt, t) loss = nn.functional.mse_loss(pred_noise, noise) optimizer.zero_grad() loss.backward() optimizer.step() losses.append(loss.item()) if step % 1000 == 0: print(f"step {step}, loss = {loss.item():.4f}") model = Denoiser() optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) T = 200 betas, alphas, alpha_bar = create_noise_schedule(T) train(model, optimizer, alpha_bar=alpha_bar)

训练超参数我实测下来比较合适的组合如下：

3.5 采样与结果可视化

训练完成后采样过程是扩散模型的精华所在。从标准正态分布随机初始化一批点，然后从T步逐步往回到第0步，每一步都用网络预测的噪声来修正当前样本：

@torch.no_grad() def sample(model, n_points=800, T=200, alphas=None, alpha_bar=None, betas=None): model.eval() x = torch.randn(n_points, 2) for t in range(T - 1, -1, -1): t_tensor = torch.full((n_points,), t, dtype=torch.float32) pred_noise = model(x, t_tensor) alpha_t = alphas[t] alpha_bar_t = alpha_bar[t] sigma_t = torch.sqrt(betas[t]) x = 1 / torch.sqrt(alpha_t) * ( x - (1 - alpha_t) / torch.sqrt(1 - alpha_bar_t) * pred_noise ) if t > 0: x = x + sigma_t * torch.randn_like(x) return x

把采样的点画出来，你会看到一堆初始杂乱的噪声点，一点点聚拢成S型。我做了一个可视化脚本，生成800个点，同时画出理论sigmoid曲线做对比：

import matplotlib.pyplot as plt samples = sample(model, n_points=800, T=T, alphas=alphas, alpha_bar=alpha_bar, betas=betas).numpy() x_grid = np.linspace(-6, 6, 300) y_theory = sigmoid(x_grid) plt.figure(figsize=(8, 5)) plt.scatter(samples[:, 0], samples[:, 1], s=6, alpha=0.6, label="generated", color="#4C72B0") plt.plot(x_grid, y_theory, "r--", linewidth=2, label="true sigmoid") plt.legend() plt.xlabel("x") plt.ylabel("y") plt.title("Diffusion Model Generated S-curve") plt.xlim([-7, 7]) plt.ylim([-0.2, 1.2]) plt.show()

我每次跑完这个可视化，都有一种“啊这就是生成模型”的实感。模型并没有见过任何真实分布公式，仅仅从加噪-去噪的监督信号中，就学会了S型曲线的形状约束，还无师自通地学会了x范围要落在-6到6之间，y要在0到1附近——这些都是从数据分布里提炼出来的。

4. 常见问题与排坑实录

4.1 训练loss不降或直接崩溃

这是最常遇到的问题。排查顺序建议从三个地方看起：先看网络里有没有加时间嵌入，不看t的网络就像蒙眼去噪，loss自然降不下去；再看学习率，1e-3是合理起点，超过5e-3容易发散；最后看beta_max，超过0.05会让最后几步噪声太猛，模型训练困难。

我见过好几次小白踩坑把T改成10试图“加速”，结果怎么都训练不好。原因很简单：T太小，每步的加噪幅度太大，闭式公式的近似程度下降，同时模型要去预测“一大团噪声”里的细节，难度剧增。对于这个demo，T低于50就会出现生成质量明显下降，建议保持200。

4.2 生成的点云又散又糊，S形状不明显

造成这种结果的原因通常是训练步数不足或者数据噪声scale设太大。训练不到1000步时，模型只学到了“大致分布在中间区域”，还没能力精确还原曲线带状结构，把步数跑到6000左右就有明显改善。

数据噪声scale也是一个关键因素，我默认是0.05，如果调到0.3，生成的点云会散得很难看。你可以做一个简单实验，分别用0.02和0.2训练，对比非常明显，能直观理解“训练数据分布宽窄对生成结果的影响”。

4.3 采样时最后一步没有加噪声控制

这是DDPM采样最容易漏的细节。如果最后一步（t=0）仍然注入随机噪声，生成点看起来会“发毛”，像是点云表面被撒了一层砂子。一定要加上if t > 0这个判断，最后一步只做均值去噪。同样的道理也适用于噪声调度器设计，理解了这个细节就理解了为什么很多采样实现里会对最后一步做特殊处理。

4.4 生成点云整体偏移或者分布范围不对

如果你观察到生成点的y轴范围跑到了比如-3到4，那大概率是数据没有做标准化。神经网络对输入输出的尺度很敏感，虽然sigmoid的y天然在0到1区间，但x范围是-6到6，两个维度的尺度差异会让网络训练变慢，甚至学到扭曲的分布。

正确的做法是把输入数据标准化到零均值、单位方差，生成后在反标准化回来。不过对这个demo来说，sigmoid函数天然把y规整到0到1，直接训练通常也能收敛。如果你改装成其他曲线（比如抛物线y=x^2、圆环等），就一定记得先做了标准化再训练。

4.5 如何判断模型真的学会了“S型分布”

除了肉眼观察，我建议你做更定量一点的验证：统计生成点在不同x区间下y的均值与标准差，和理论sigmoid曲线比较；或者直接用RMSE评估生成点云与真实曲线的接近程度。具体做法是把x轴分桶，计算每个桶内生成点的y均值，再求与真实sigmoid值的均方根误差，误差小于0.1基本就算学会S型分布。

另外一个更“生成模型”角度的验证方法：用不同随机种子采样多批点，看每批点云的整体形态是否稳定。如果每批都是完整S型，说明模型学到了真实的分布，而不是仅仅记住了某个固定形状。

写在最后：跑完demo之后建议做的三件事

我个人跑完这个demo后最大的感受是：扩散模型没有想象中那么“玄学”。它本质上就是一个加了时间条件的去噪网络，训练目标就是预测噪声，采样过程就是从噪声中逐步“剥离”出数据。这个认知帮我建立起了后续学习所有扩散模型变体的底层框架。

最后再分享三个我反复试验过的扩展玩法，都不需要原始代码大改：第一，把数据生成函数从sigmoid换成圆形分布（angle均匀采样，半径加噪声），训练出来的点云就是一个甜甜圈，这个过程你会更深刻地理解条件分布和流形学习；第二，把T从200改成1000，对比一下生成质量和耗时，你会直观感受到步数对生成精度的权衡；第三，把中间采样的每一帧都保存下来做成GIF动画，看着随机噪声一点点聚拢成S型的全过程，那种感觉比自己画一百张曲线图都来得直观。等你能独立完成这三个实验，再回去看DDPM原论文，你会发现自己理解速度怕是要快上一倍不止。

本文还有配套的精品资源，点击获取