| import torch | |
| # ========================= | |
| # Device | |
| # ========================= | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| # ========================= | |
| # Dataset | |
| # ========================= | |
| IMAGE_SIZE = 28 | |
| CHANNELS = 1 | |
| # ========================= | |
| # DDPM | |
| # ========================= | |
| TIMESTEPS = 1000 | |
| BETA_START = 1e-4 | |
| BETA_END = 0.02 | |
| # ========================= | |
| # Model | |
| # ========================= | |
| TIME_EMBED_DIM = 128 | |
| BASE_CHANNELS = 64 | |
| MID_CHANNELS = 128 | |
| # ========================= | |
| # Training | |
| # ========================= | |
| BATCH_SIZE = 128 | |
| EPOCHS = 20 | |
| LEARNING_RATE = 2e-4 | |
| WEIGHT_DECAY = 1e-4 | |
| # ========================= | |
| # Checkpoints | |
| # ========================= | |
| CHECKPOINT_DIR = "checkpoints" | |
| MODEL_NAME = "ddpm_fashion_mnist.safetensors" | |
| # ========================= | |
| # Sampling | |
| # ========================= | |
| NUM_SAMPLES = 16 | |
| # ========================= | |
| # Random Seed | |
| # ========================= | |
| SEED = 42 | |