# stable_diffusion_mnist

**Repository Path**: giteepro/stable_diffusion_mnist

## Basic Information

- **Project Name**: stable_diffusion_mnist
- **Description**: 使用mnist数据集理解DDPM，代码基于网络修改。
- **Primary Language**: Unknown
- **License**: Not specified
- **Default Branch**: master
- **Homepage**: None
- **GVP Project**: No

## Statistics

- **Stars**: 0
- **Forks**: 0
- **Created**: 2026-04-12
- **Last Updated**: 2026-05-05

## Categories & Tags

**Categories**: Uncategorized

**Tags**: None

## README

# Diffusion Model - DDPM on MNIST

基于去噪扩散概率模型（DDPM）的手写数字生成实现。

---

## 项目文件结构

```
diffusion-model/
├── generate_standard_mnist.py    # 生成标准MNIST数据集（JPG格式）
├── main_images.py                 # DDPM训练与测试脚本（使用JPG图像）
├── verify_data.py                 # 数据校验脚本
├── u_net.pt                       # 训练好的U-Net模型权重
├── dataset/
│   ├── train/                     # 训练集图像（60000张JPG）
│   └── test/                      # 测试集图像（10000张JPG）
└── data/                          # torchvision下载的原始MNIST数据
```

---

## DDPM 理论核心

DDPM 包含两个过程：

### 1. 前向扩散过程（Forward Diffusion）
- 逐步向真实数据添加高斯噪声
- x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon
- 其中 alpha_bar_t = 从 i=1 到 t 的连乘 (1 - beta_i)

### 2. 反向去噪过程（Reverse Denoising）
- 学习从噪声逐步恢复真实数据
- 使用 U-Net 预测噪声 epsilon_theta
- 通过 x_{t-1} ~ p_theta(x_{t-1}|x_t) 采样

---

## 一、数据集生成

### 生成标准MNIST数据集

使用 `generate_standard_mnist.py` 从 torchvision 下载标准MNIST数据集并保存为JPG格式：

```bash
python generate_standard_mnist.py
```

**功能说明：**
- 自动下载MNIST数据集（训练集60000张，测试集10000张）
- 将图像保存为JPG格式到 `dataset/train/` 和 `dataset/test/`
- 图像命名格式：`{label}_{index:06d}.jpg`，例如 `5_000001.jpg` 表示标签为5的第2张图像
- 图像自动归一化到 [0, 255] 范围

---

## 二、数据校验

使用 `verify_data.py` 验证数据集的正确性：

```bash
python verify_data.py
```

**校验内容：**
- 训练集和测试集的图像数量
- 图像形状（应为784维向量，即28×28）
- 数据范围（归一化到 [0, 1]）
- 标签分布（0-9每个数字的样本数量）
- 标签范围（0-9）

---

## 三、数据读取

### `read_images_from_directory()` 函数

`main_images.py` 中的数据读取函数使用 OpenCV 读取JPG图像：

```python
def read_images_from_directory(dir_path):
    """从目录中读取所有图像和标签（使用OpenCV）
    返回：(图像张量, 标签张量)
    图像形状：(N, 784)，数值范围 [0, 1]
    标签形状：(N,)，数值范围 0-9
    """
```

**读取流程：**
1. 获取目录下所有 `.jpg` 文件并排序
2. 从文件名中提取标签（通过下划线分割）
3. 使用 OpenCV 以灰度模式读取图像
4. 将图像展平为784维向量
5. 归一化到 [0, 1] 范围并转换为 PyTorch 张量

---

## 四、代码模块详解

### 整体架构

**主要模块**：
1. **数据读取** - `read_images_from_directory()` 函数
2. **扩散模型核心** - `DenoiseDiffusion` 类
3. **U-Net 网络架构** - `UNet` 类及组件
4. **训练与采样** - `main()` 函数

---

### 扩散模型 (`DenoiseDiffusion`)

| 方法 | 功能 |
|------|------|
| `q_xt_x0()` | 计算给定 x_0 时 x_t 的均值和方差 |
| `q_sample()` | 从 x_0 采样得到 x_t（前向扩散） |
| `p_sample()` | 从 x_t 采样得到 x_{t-1}（反向去噪） |
| `loss()` | 计算预测噪声与真实噪声的 MSE 损失 |

**超参数**：
- beta_t: 从 0.0001 到 0.02 线性变化
- 步数 T = 1000

---

### U-Net 架构

这是一个带时间嵌入的 U-Net，用于预测噪声：

**组件**：
- `TimeEmbedding`: 正弦位置编码 + MLP，将时间步 t 映射为向量
- `ResidualBlock`: 残差块，融合时间嵌入
- `AttentionBlock`: 自注意力机制
- `DownBlock`/`UpBlock`: 下采样/上采样模块
- `MiddleBlock`: 中间瓶颈层

**架构流程**：
```
输入 x → 初始卷积 → 编码器(下采样) → 中间层 → 解码器(上采样 + 跳跃连接) → 输出噪声预测
```

---

## 五、训练流程 (`main`)

### 1. 数据准备
- 读取训练集（60000张）和测试集（10000张）
- 将训练集按 8:2 划分为训练/验证集

### 2. 模型初始化
```python
u_net = UNet(1, 16, [1, 2, 2], [False, False, False], n_blocks=1)
dm = DenoiseDiffusion(u_net, 1000, device=device)
opt_dm = torch.optim.Adam(u_net.parameters(), lr=0.001)
```

### 3. 训练循环
- 随机采样时间步 t
- 对 x_0 加噪得到 x_t
- 用 U-Net 预测噪声
- 计算 MSE 损失并反向传播
- **早停机制**：验证损失 40 轮不下降则停止训练
- 保存最佳模型到 `u_net.pt`

### 4. 采样生成
- 从纯高斯噪声 x_T ~ N(0, I) 开始
- 逐步反向采样 1000 步得到 x_0
- 每 100 步保存一次中间结果
- 最终生成 `diffusion_results.png` 展示去噪过程

---

## 六、关键数学公式对应

| 公式 | 代码位置 |
|------|----------|
| q(x_t|x_0) = N(x_t; sqrt(alpha_bar_t) * x_0, (1 - alpha_bar_t) * I) | `q_xt_x0()` |
| x_t = sqrt(alpha_bar_t) * x_0 + sqrt(1 - alpha_bar_t) * epsilon | `q_sample()` |
| mu_theta(x_t, t) = 1/sqrt(alpha_t) * (x_t - beta_t/sqrt(1 - alpha_bar_t) * epsilon_theta(x_t, t)) | `p_sample()` |
| L = E_{t,x0,epsilon}[ ||epsilon - epsilon_theta(sqrt(alpha_bar_t)*x0 + sqrt(1-alpha_bar_t)*epsilon, t)||^2 ] | `loss()` |

---

## 七、快速开始指南

### 完整运行流程

1. **生成数据集**
```bash
python generate_standard_mnist.py
```

2. **验证数据**（可选）
```bash
python verify_data.py
```

3. **训练和测试**
```bash
python main_images.py
```

### 预期输出

- 训练过程中会显示每个epoch的训练损失和验证损失
- 早停触发后会自动停止训练
- 最终生成 `diffusion_results.png`，展示从噪声逐步去噪生成手写数字的过程
- 最佳模型权重保存为 `u_net.pt`