tiny-diffusion

一个针对二维数据集的概率扩散模型的极简 PyTorch 实现。可通过运行 python ddpm.py -h 来查看训练时可用的选项，从而开始使用。

前向过程

对一千个二维点的数据集应用前向扩散过程的可视化展示。请注意，这只恐龙并不是单个训练样本，它代表了数据集中每一个二维点。

反向过程

此图展示了反向过程如何恢复训练数据的分布。

消融实验

我针对学习率、模型大小等超参数进行了一系列消融实验，并对学习过程进行了可视化。图表中的列代表检查点的轮次，行则表示超参数的取值。每个单元格显示了一千个生成的二维点。

学习率

学习过程对学习率非常敏感。起初，模型的输出效果很差，让我一度怀疑代码中存在错误。然而，仅仅调整学习率的数值便解决了问题。

数据集

当前的模型配置在 line 数据集上表现不佳，而我认为这是所有数据集中最基础的一个。理想情况下，图像的角应该是清晰锐利的，但实际上却显得模糊。

时间步数

扩散过程的时间步数越多，生成的结果越好。如果时间步数较少，则生成的“恐龙”会不完整，顶部和底部都会缺失一些点。

方差调度

二次方调度并未带来更好的效果。未来还可以尝试余弦或 S 形等其他调度方式。

隐藏层大小

模型的容量似乎并不是瓶颈，因为在不同隐藏层大小下都能得到相似的结果。

隐藏层数量

与隐藏层大小的消融实验类似，模型的容量似乎也并非限制因素。

位置嵌入（时间步）

模型确实受益于时间步信息，但具体采用何种编码方式并不重要。

位置嵌入（输入）

对于低维问题域中的高频函数学习，使用正弦嵌入来处理输入是有帮助的。例如，在 这项研究 中，就展示了如何将每个 (x, y) 像素坐标映射到 (r, g, b) 颜色。同样的道理也适用于当前场景。

参考文献

• 恐龙数据集来源于 Datasaurus Dozen 数据。
• HuggingFace 的 diffusers 库。
• lucidrains 的 PyTorch 版 DDPM 实现。
• Jonathan Ho 的 DDPM 实现。
• InFoCusp 的 TensorFlow 版 DDPM 实现。

tiny-diffusion 快速上手指南

tiny-diffusion 是一个极简的 PyTorch 实现，专为 2D 数据集的概率扩散模型设计。本指南将帮助你快速配置环境并运行训练。

环境准备

在开始之前，请确保你的开发环境满足以下要求：

• 操作系统：Linux、macOS 或 Windows
• Python 版本：推荐 Python 3.8 及以上
• 核心依赖：
  • PyTorch (需包含 CUDA 支持以加速训练，可选 CPU 模式)
  • 其他标准科学计算库（如 numpy, matplotlib 等，通常随 PyTorch 生态安装）

国内加速建议： 推荐使用清华或阿里镜像源安装 PyTorch，以提升下载速度。

• 清华源：https://pypi.tuna.tsinghua.edu.cn/simple
• PyTorch 官方中文镜像：https://download.pytorch.org/whl/torch_stable.html

安装步骤

1. 克隆项目仓库

   git clone https://github.com/your-username/tiny-diffusion.git
   cd tiny-diffusion
   

   (注：请将 URL 替换为实际的项目仓库地址)

2. 安装依赖 建议使用虚拟环境（如 venv 或 conda）以避免依赖冲突。

   使用 pip 安装（推荐国内镜像）：

   pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   若项目中无 requirements.txt，手动安装核心依赖：

   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
   pip install numpy matplotlib -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   (注：cu118 代表 CUDA 11.8，请根据你的显卡驱动版本调整，若无 GPU 可移除 --index-url 参数或使用 CPU 版本)

基本使用

该项目最核心的功能是训练和可视化 2D 点云的扩散过程。

1. 查看帮助信息 运行以下命令以探索所有可用的训练选项和参数：

   python ddpm.py -h
   

2. 开始训练 使用默认配置启动训练（通常会自动加载内置的 "dino" 2D 点云数据集）：

   python ddpm.py
   

3. 自定义超参数（示例） 根据 README 中的消融实验结论，你可以调整学习率或时间步数来优化结果。例如，设置更大的时间步数以获得更完整的生成效果：

   python ddpm.py --num_timesteps 1000 --learning_rate 1e-4
   

训练完成后，脚本通常会生成可视化图像（如前向扩散过程和反向生成过程的分布图），展示模型如何从噪声中恢复出原始数据的分布（如恐龙形状的点云）。