blob: b4f2756f3b09f1865fdbed74ab375bdf0922847e (
plain) (
blame)
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
52
53
54
55
|
# MNIST Diffusion Model
A diffusion-based generative model for MNIST digits implemented in PyTorch.
## Installation
```bash
pip install -r requirements.txt
```
## Usage
### Train Diffusion Model
```bash
python main.py train --epochs 10 --lr 2e-4 --batch-size 64
```
### Train with Self-Attention
```bash
python main.py train --epochs 10 --attention
```
### Train Autoencoder (for latent diffusion)
```bash
python main.py train-ae --epochs 10
```
### Generate Samples
```bash
python main.py sample --checkpoint checkpoints/diffusion_latest.pt
```
### Visualize Diffusion Process
```bash
python main.py visualize --all
```
## Configuration
All hyperparameters can be found in `src/config.py`:
- `DIFFU_STEPS`: Number of diffusion steps (default: 1000)
- `EPOCHS`: Training epochs (default: 10)
- `BATCH_SIZE`: Batch size (default: 64)
- `LEARNING_RATE`: Learning rate (default: 2e-4)
## Model Architecture
The diffusion model uses a UNet architecture with:
- Timestep embedding
- Label conditioning (for class-conditional generation)
- Optional self-attention layers
## License
See [LICENSE](LICENSE) for details.
|
