DDPM 学习 Demo（PyTorch + Conda + GPU）

这是一个最小可学习的 DDPM 示例，面向入门理解：

• 正向扩散：逐步给图像加噪
• 反向去噪：模型学习预测噪声
• 训练目标：$\mathcal{L}=\mathbb{E}[|\epsilon-\epsilon_\theta(x_t,t)|^2]$

数据集默认使用 MNIST，训练后会周期性采样并输出图片。

当前版本新增：

• CUDA GPU 加速（适配 RTX 4060）
• AMP 混合精度训练（默认自动开启）
• EMA（指数滑动平均）提升采样质量
• Cosine 噪声调度（默认）
• 按时间步 t 的加噪轨迹可视化
• 按时间步 t 的去噪轨迹可视化

1. 创建 conda 环境

conda env create -f environment.yml
conda activate ddpm-demo

如果环境已存在，使用：

conda env update -f environment.yml --prune

验证 GPU 是否可用：

python -c "import torch; print(torch.cuda.is_available(), torch.cuda.get_device_name(0) if torch.cuda.is_available() else 'CPU')"

2. 开始训练

python ddpm_demo.py --mode train --epochs 12 --batch-size 256 --timesteps 300 --schedule cosine --variance-type posterior --base-ch 96 --amp --sample-every 1000 --save-dir outputs_better

上面这组参数是当前仓库里已实测、去噪效果更稳定的一组推荐配置（RTX 4060）。

常用参数：

• --epochs：训练轮数
• --timesteps：扩散步数（越大质量通常更好，但更慢）
• --schedule：噪声调度（cosine 或 linear，默认 cosine）
• --variance-type：反向采样方差（beta 或 posterior，默认 beta 以兼容旧 checkpoint）
• --amp：开启混合精度（CUDA 上更快）
• --ema-decay：EMA 衰减系数（默认 0.9999）
• --sample-every：每多少 step 采样一次
• --noise-vis-steps：保存加噪轨迹图时的 t 采样点数量
• --denoise-vis-steps：保存去噪轨迹图时的 t 采样点数量
• --save-dir：模型和样本输出目录
• --ignore-ckpt-config：采样时忽略 checkpoint 内训练参数（默认不建议开启）
• --use-raw-model：采样时使用原始模型权重（短训练下可能比 EMA 更好）

训练输出：

• 模型权重：outputs/ddpm_epoch_*.pt
• 过程采样图：outputs/sample_step_*.png
• 加噪轨迹图：outputs/noise_trajectory.png
• 去噪轨迹图：outputs/denoise_trajectory.png

3. 从 checkpoint 采样

python ddpm_demo.py --mode sample --ckpt outputs_better/ddpm_epoch_12.pt --num-sample-images 16 --save-dir outputs_better --sample-name sample_best.png --use-raw-model

说明：默认会自动读取 checkpoint 中保存的 timesteps / schedule / variance-type，避免参数错配导致去噪失败。 仅在你明确要做对比实验时，才建议加 --ignore-ckpt-config 强制使用命令行参数。

将生成：outputs/sample_from_ckpt.png

示例命令中会生成：outputs_better/sample_best.png

4. 单独生成“每个 t 的加噪过程图”

python ddpm_demo.py --mode visualize-noise --timesteps 300 --noise-vis-steps 12

将生成：outputs/noise_trajectory.png

5. 单独生成“每个 t 的去噪过程图”

python ddpm_demo.py --mode visualize-denoise --ckpt outputs_better/ddpm_epoch_12.pt --denoise-vis-steps 12 --save-dir outputs_better --denoise-vis-name denoise_best.png --use-raw-model

将生成：outputs_better/denoise_best.png

也可以在普通采样时一起保存去噪轨迹：

python ddpm_demo.py --mode sample --ckpt outputs_better/ddpm_epoch_12.pt --save-dir outputs_better --sample-name sample_with_denoise.png --denoise-vis-steps 12 --denoise-vis-name denoise_from_sample.png --use-raw-model

如果使用“旧版本脚本训练出来的 checkpoint”，建议固定：

• --schedule linear
• --variance-type beta

6. 推荐学习路径

1. 先把 --timesteps 改成 100 快速验证流程。
2. 再改回 300 或 1000 观察采样质量变化。
3. 尝试增加 --base-ch（比如 128）对比效果与速度。
4. 阅读 ddpm_demo.py 中 q_sample / p_sample 对应公式实现。

7. 注意事项

• 首次运行会自动下载 MNIST 到 data/。
• 如果显存不足，减小 --batch-size。
• 如果只用 CPU，训练会较慢，可先减少 --epochs 和 --timesteps。
• 在 RTX 4060 上，建议优先使用 --amp、--batch-size 256（或按显存调整）与 --schedule cosine。
• 如果“去噪效果很差”，优先排查：
• 训练轮数是否太少（建议先提升到 20+）
• 是否误用与 checkpoint 不一致的 timesteps/schedule/variance-type
• 是否使用了旧 checkpoint（可先重训 1 版再对比）
• 若短轮数训练且结果偏噪声，可尝试 --use-raw-model 跳过 EMA 进行采样
• 当前实测中，12 轮时 --use-raw-model 比 EMA 采样更清晰；训练更久后可再对比 EMA。