Function train

skills/paper2code/worked/ddpm/src/train.py:43–186

Algorithm 1 — DDPM Training. Args: config_path: Path to YAML config file

(config_path: str = "configs/base.yaml")

def train(config_path: str = "configs/base.yaml"):
    """Algorithm 1 — DDPM Training.

    Args:
        config_path: Path to YAML config file
    """
    # --- Load config ---
    config_path = Path(config_path)
    if config_path.exists():
        with open(config_path) as f:
            cfg = yaml.safe_load(f)
    else:
        raise FileNotFoundError(f"Config not found: {config_path}")

    diff_cfg = cfg["diffusion"]
    model_cfg = cfg["model"]
    train_cfg = cfg["training"]
    data_cfg = cfg["data"]

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Using device: {device}")

    # --- Noise schedule ---
    # §2, Eq. 4 — linear schedule β_1 = 0.0001, β_T = 0.02
    T = diff_cfg["T"]
    betas = linear_noise_schedule(T, diff_cfg["beta_start"], diff_cfg["beta_end"])
    betas = betas.to(device)

    alphas = 1.0 - betas                                    # α_t = 1 − β_t
    alpha_bar = torch.cumprod(alphas, dim=0)                # ᾱ_t = ∏_{s=1}^{t} α_s
    sqrt_alpha_bar = torch.sqrt(alpha_bar)                  # √ᾱ_t
    sqrt_one_minus_alpha_bar = torch.sqrt(1.0 - alpha_bar)  # √(1−ᾱ_t)

    # --- Model ---
    unet_config = UNetConfig(
        image_channels=model_cfg.get("image_channels", 3),
        base_channels=model_cfg.get("base_channels", 128),
        channel_mults=tuple(model_cfg.get("channel_mults", [1, 2, 2, 2])),
        num_res_blocks=model_cfg.get("num_res_blocks", 2),
        attention_resolutions=tuple(model_cfg.get("attention_resolutions", [16])),
        dropout=model_cfg.get("dropout", 0.0),
        num_groups=model_cfg.get("num_groups", 32),
        image_size=data_cfg.get("image_size", 32),
    )
    model = UNet(unet_config).to(device)
    logger.info(f"Model: {model}")

    # --- EMA ---
    # §4 — "we also report results with an exponential moving average of
    # model parameters with a decay factor of 0.9999"
    ema = EMA(model, decay=train_cfg.get("ema_decay", 0.9999))

    # --- Optimizer ---
    # Appendix B — "Adam, lr = 2 × 10^-4"
    optimizer = torch.optim.Adam(
        model.parameters(),
        lr=float(train_cfg.get("lr", 2e-4)),
    )