File size: 1,881 Bytes
96a0bf1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 | ---
license: mit
pipeline_tag: unconditional-image-generation
tags:
- diffusion
- ddpm
- cifar10
- onnx
- webgpu
datasets:
- uoft-cs/cifar10
---
# CIFAR-10 DDPM (18.7M params, from scratch)
A small denoising diffusion probabilistic model ([DDPM, Ho et al. 2020](https://arxiv.org/abs/2006.11239)) trained on CIFAR-10, written from scratch in ~250 lines of plain PyTorch.
**[Try it in your browser](https://huggingface.co/spaces/PeterRabbit/cifar10-ddpm-demo)** — the ONNX export runs client-side on WebGPU, generating an image in ~2 seconds on your own GPU.
*Uncurated 8×8 grid after 300 epochs.*
## Model details
- U-Net, **18.7M parameters**: three resolution levels (32→16→8), residual blocks with sinusoidal timestep embeddings, self-attention at 16×16 and 8×8, dropout 0.1
- Linear beta schedule, T=1000, trained with the simplified noise-prediction MSE objective
- EMA of weights (decay 0.9995) used for sampling; DDIM sampling with 50 steps
- Unconditional — no class or text control
- Trained 300 epochs (~117k steps, batch 128) in ~4 hours on a single RTX 3090; final loss ≈ 0.029
## Files
| File | Description |
|---|---|
| `checkpoint.pt` | Full PyTorch checkpoint: raw weights, EMA weights, optimizer state (resumable) |
| `unet.onnx` | EMA weights exported to ONNX (fp32), verified to 4e-6 against PyTorch |
| `train_diffusion.py` | Complete training script |
| `sample.py` | Generate image grids from the checkpoint |
| `export_onnx.py` | Reproduce the ONNX export |
## Usage
```bash
pip install torch torchvision
python sample.py --n 64 --seed 42 # sample a grid from checkpoint.pt
python train_diffusion.py --base 128 --epochs 300 --out out_big --sample-every 10 --ema-decay 0.9995 # retrain
```
Source and local web app: [github.com/dannysheesh/cifar10-ddpm](https://github.com/dannysheesh/cifar10-ddpm)
|