| """Method interface. EDM, Cascaded SI and Direct SI subclass this.""" |
|
|
| from abc import ABC, abstractmethod |
|
|
| import torch |
|
|
|
|
| class Method(ABC): |
| """Training-method interface. |
| |
| A method bundles together (a) the training-time loss, (b) the sampling |
| procedure used at eval time, and (c) the optimizer for the model. The |
| trainer is method-agnostic — it calls these three hooks and nothing else. |
| """ |
|
|
| @abstractmethod |
| def loss(self, model: torch.nn.Module, batch: tuple[torch.Tensor, ...]) -> torch.Tensor: |
| """Per-batch training loss. batch = (x_0, x_T) tensors.""" |
|
|
| @abstractmethod |
| def sample(self, model: torch.nn.Module, x_0: torch.Tensor, n_samples: int, |
| **variant_kwargs) -> torch.Tensor: |
| """Generate n_samples per x_0 conditioning vector. Returns [B, n_samples, D]. |
| |
| Methods that support multiple eval variants accept variant-specific |
| overrides as keyword arguments (e.g. ``noise_scale=0.25`` for SI). The |
| legal kwargs are whatever ``variants_for_eval()`` produces. |
| """ |
|
|
| def variants_for_eval(self) -> dict[str, dict]: |
| """Eval-time sampling variants this method exposes. |
| |
| Returns a dict ``{variant_label: sample_kwargs}``; ``bifurcation.sample`` writes |
| one directory per variant. |
| |
| Default: one variant called ``default`` with no kwarg overrides — used |
| by EDM (the Heun sampler is deterministic; there's no analogous |
| noise-scale knob). |
| """ |
| return {"default": {}} |
|
|
| def configure_optimizer(self, model: torch.nn.Module, cfg: dict) -> torch.optim.Optimizer: |
| """Build the optimizer from ``self.optimizer_cfg`` (set by the method's __init__). |
| |
| Supported types: |
| - ``adam`` (default) → ``torch.optim.Adam`` ; weight_decay ignored |
| - ``adamw`` → ``torch.optim.AdamW``; ``weight_decay`` honoured |
| (true decoupled decay — the right way to regularise |
| adaptive-LR optimisers) |
| |
| Default (no ``optimizer`` block in the method section, or |
| ``optimizer.type == "adam"`` with ``weight_decay`` 0): plain Adam. |
| |
| A non-None ``cfg["weight_decay"]`` (the training section, passed by |
| the trainer) overrides the method's block and selects AdamW with that decay. |
| """ |
| opt_cfg = getattr(self, "optimizer_cfg", None) or {} |
| opt_type = str(opt_cfg.get("type", "adam")).lower() |
| wd = float(opt_cfg.get("weight_decay", 0.0)) |
| params = {"lr": cfg["lr"], "betas": (0.9, 0.999), "eps": 1e-8} |
| wd_override = cfg.get("weight_decay") |
| if wd_override is not None: |
| return torch.optim.AdamW(model.parameters(), |
| weight_decay=float(wd_override), **params) |
| if opt_type == "adam": |
| return torch.optim.Adam(model.parameters(), **params) |
| if opt_type == "adamw": |
| return torch.optim.AdamW(model.parameters(), weight_decay=wd, **params) |
| raise ValueError(f"unknown optimizer.type: {opt_type!r}") |
|
|
