File size: 3,117 Bytes
41dd7e0 | 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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """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}")
|