NJU R&L Group Assessment: Generative AI

Checkpoints are available at Google Drive Link

1. 项目概述

核心目标：基于 CIFAR-10 数据集，分别实现并训练两种主流生成模型架构，并对比其优劣。

核心任务

• Task 1 (Diffusion): 训练一个基于 U-Net 架构的 Diffusion Model，实现图像随机生成。
• Task 2 (Autoregressive): 训练一个基于 GPT 架构的 Autoregressive Model，实现图像随机生成。
• Task 3 & 4 (Analysis): 对比两种模型的优劣，探讨生成模型与人类创作的异同。

---

2. 环境配置

本项目在服务器上使用 Docker 进行开发。

Docker 配置

• 基础镜像: pytorch/pytorch:2.4.0-cuda12.4-cudnn9-devel
• 容器管理: 使用 docker compose 进行一键式启动。

快速开始

1. 构建并启动容器:

   docker compose up -d --build

2. 进入容器环境:

   docker exec -it genai_container bash

3. 模型训练:
   • Task 1: Diffusion Model

     python train.py --config configs/diffusion_unet.yaml

   • Task 2: VQ-VAE + GPT Step 1: 训练 VQ-VAE

     python train.py --config configs/vqvae.yaml

     Step 2: 训练 GPT (AR)

     python train.py --config configs/vqvae_gpt.yaml

快速采样 (Sampling)

训练完成后，你可以使用 sample.py 脚本生成图像。

1. Diffusion 模型采样

• DDPM (1000 steps):

  python sample.py 
      --config configs/diffusion_unet.yaml 
      --task diffusion 
      --scheduler ddpm 
      --checkpoint outputs/diffusion/model_epoch_700.pth 
      --num_samples 16 
      --output_path outputs/diffusion_samples_ddpm.png

• DDIM (50 steps, 快速采样):

  python sample.py 
      --config configs/diffusion_unet.yaml 
      --task diffusion 
      --scheduler ddim 
      --sample_steps 50 
      --checkpoint outputs/diffusion/model_epoch_700.pth 
      --num_samples 16 
      --output_path outputs/diffusion_samples_ddim.png

2. Autoregressive (VQ-VAE + GPT) 采样

python sample.py 
    --config configs/vqvae_gpt.yaml 
    --task ar 
    --checkpoint outputs/vqvae_gpt/model_epoch_4000.pth 
    --vqvae_ckpt outputs/vqvae/model_epoch_600.pth 
    --num_samples 16 
    --top_k 50 
    --temp 1.0 
    --output_path outputs/ar_samples.png

---

3. 代码架构 (Code Architecture)

本项目采用高度模块化、配置驱动的设计模式。

目录结构树

GenAI_Project/
├── configs/                   # [Configuration] 所有超参数的 YAML 文件
│   ├── diffusion_unet.yaml    # Task 1: Diffusion 配置
│   ├── vqvae.yaml             # Task 2: VQ-VAE 配置
│   └── vqvae_gpt.yaml         # Task 2: GPT (AR) 配置
├── data/                      # [Data] 原始数据集与预处理数据目录
├── outputs/                   # [Outputs] 训练产物目录 (模型权重、生成样例、训练日志)
├── src/                       # [Source Code] 核心代码库
│   ├── data/                  # 数据管道
│   │   └── cifar_loader.py    # 定义 Dataset, DataLoader 与 数据增强
│   ├── models/                # 模型定义
│   │   ├── modules.py         # 通用底层组件 (RMSNorm, SwiGLU, FlashAttention)
│   │   ├── vqvae.py           # VQ-VAE 模型实现 (Encoder/Decoder/Quantizer)
│   │   ├── diffusion/         # Task 1 核心模块
│   │   │   ├── unet.py        
│   │   │   └── scheduler.py   
│   │   └── autoregressive/    # Task 2 核心模块
│   │       └── gpt.py         
│   ├── engine/                # 训练与推理流程控制 
│   │   ├── trainer.py         # [Base] 基础训练器 (负责设备管理, Checkpoint, 指标记录)
│   │   ├── vqvae_trainer.py   # [VQ-VAE] 实现 Reconstruction & Commitment Loss
│   │   ├── diffusion_trainer.py # [Diffusion] 实现 Noise Prediction MSE Loss
│   │   └── ar_trainer.py      # [GPT] 实现 Next Token Prediction 交叉熵 Loss
│   └── utils/                 # 工具箱
│       ├── config_parser.py   # 动态 YAML 配置加载器
│       ├── logger.py          # 终端输出与实验日志管理
│       ├── visualizer.py      # 生成图像的网格化保存工具
│       ├── plot_diffusion.py  # Diffusion 训练过程可视化
│       └── plot_gpt.py        # GPT 训练过程可视化
├── train.py                   # [Entry Point] 统一训练入口 (多任务路由)
├── sample.py                  # [Entry Point] 统一采样入口 (支持 DDPM/DDIM/AR)
├── Dockerfile                 # 环境构建镜像定义
├── docker-compose.yml         # 快捷容器部署配置
├── requirements.txt           # Python 环境依赖包
└── README.md                  # 本项目技术文档

---

4. 性能表现 (Performance & Benchmarks)

在服务器环境下生成 16 张图片 (Batch Size = 16) 的耗时对比：

模型方案	采样器	步数 (Steps)	耗时 (Latency)	效率 (it/s)
Diffusion	DDPM	1000	~11.9s	~84 it/s
Diffusion	DDIM	50	~0.95s	~52 it/s
AR (GPT)	Top-K	64	~0.60s	~106 it/s

注：AR 模型在生成离散 Token 时具有天然的速度优势，而 DDIM 极大地提升了扩散模型的采样效率。

---

5. 模块详细职责 (Module Responsibilities)

A. 配置层 (configs/)

• 原则：代码中严禁硬编码超参数（Learning Rate, Layers, Batch Size 等）。
• 格式：YAML。
• 训练配置:
  • diffusion_unet.yaml: 扩散步数、Beta 调度、U-Net 通道数。
  • vqvae.yaml: Codebook 大小、压缩率。
  • vqvae_gpt.yaml: GPT 层数、注意力头数、Patch 尺寸。

B. 数据层 (src/data/cifar_loader.py)

• 功能：标准化 CIFAR-10 数据流入。
• 预处理区分：
  • Diffusion: [-1, 1] 归一化。
  • GPT: 转换为离散 Token 索引序列 (基于 VQ-VAE)。

C. 模型层 (src/models/)

通用组件 (modules.py)

• 实现高性能可复用层：RMSNorm, SwiGLU , Scaled Dot-Product Attention (Flash Attention)。

Task 1: Diffusion (diffusion/)

• unet.py: 实现带时间步注入的卷积网络。
• scheduler.py: 封装扩散数学。支持 DDPM (随机性) 与 DDIM (确定性/跳步采样)。

Task 2: Autoregressive (autoregressive/ & vqvae.py)

• gpt.py: 纯 Decoder 结构的 Transformer，用于预测下一个 Token 索引。
• vqvae.py: 基于 EMA 更新的向量量化自编码器，将 $32 \times 32 \times 3$ 图像压缩为 $8 \times 8$ 的离散 Latent。

D. 引擎层与工具类 (src/engine/ & src/utils/)

本项目使用 Template Method Pattern (模板方法模式) 极大地提升了代码复用率：

• trainer.py (BaseTrainer): 抽象了模型保存、优化器加载、通用的训练循环。
• diffusion_trainer.py / vqvae_trainer.py / ar_trainer.py: 仅需通过实现 train_step 钩子函数，即可定制不同生成范式的 Loss 计算。

---