Function load_model

skills/paper2code/worked/ddpm/src/evaluate.py:41–87

Load a trained DDPM model from checkpoint. §4 — "we also report results with an exponential moving average" The EMA parameters typically produce better samples. Args: checkpoint_path: Path to .pt checkpoint file device: Target device use_ema: Whether to load EMA

(
    checkpoint_path: str,
    device: torch.device,
    use_ema: bool = True,
)

def load_model(
    checkpoint_path: str,
    device: torch.device,
    use_ema: bool = True,
) -> tuple:
    """Load a trained DDPM model from checkpoint.

    §4 — "we also report results with an exponential moving average"
    The EMA parameters typically produce better samples.

    Args:
        checkpoint_path: Path to .pt checkpoint file
        device: Target device
        use_ema: Whether to load EMA parameters (recommended)

    Returns:
        (model, config_dict)
    """
    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
    cfg = checkpoint["config"]
    model_cfg = cfg["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=cfg["data"].get("image_size", 32),
    )

    model = UNet(unet_config).to(device)

    if use_ema and "ema_state_dict" in checkpoint:
        # Load EMA parameters
        ema_params = checkpoint["ema_state_dict"]
        for name, param in model.named_parameters():
            if name in ema_params:
                param.data.copy_(ema_params[name])
        logger.info("Loaded EMA parameters")
    else:
        model.load_state_dict(checkpoint["model_state_dict"])
        logger.info("Loaded model parameters (no EMA)")

    return model, cfg


@torch.no_grad()