Training

A generic teacher-training loop with checkpoint I/O.

polyweave.training

Training utilities: a generic teacher-training loop and checkpoint I/O.

TeacherTrainResult dataclass

TeacherTrainResult(losses: List[float] = list(), pi_scales: List[float] = list(), exponent_abs_means: List[float] = list(), pi_shares: List[float] = list())

Outcome of :func:train_teacher.

load_checkpoint

load_checkpoint(path: str, model: Module, optimizer: Optional[Optimizer] = None, map_location: Optional[Any] = None, strict: bool = True) -> Dict[str, Any]

Load weights into model (and optionally optimizer) from path.

Returns the stored meta dict (empty if none was saved).

Source code in polyweave/training/checkpoint.py

def load_checkpoint(
    path: str,
    model: nn.Module,
    optimizer: Optional[torch.optim.Optimizer] = None,
    map_location: Optional[Any] = None,
    strict: bool = True,
) -> Dict[str, Any]:
    """Load weights into ``model`` (and optionally ``optimizer``) from ``path``.

    Returns the stored ``meta`` dict (empty if none was saved).
    """
    payload = torch.load(path, map_location=map_location, weights_only=False)
    model.load_state_dict(payload["model_state"], strict=strict)
    if optimizer is not None and "optimizer_state" in payload:
        optimizer.load_state_dict(payload["optimizer_state"])
    return payload.get("meta", {})

save_checkpoint

save_checkpoint(path: str, model: Module, optimizer: Optional[Optimizer] = None, meta: Optional[Dict[str, Any]] = None) -> None

Save model (and optionally optimizer/meta) to path.

Parent directories are created if needed.

Source code in polyweave/training/checkpoint.py

def save_checkpoint(
    path: str,
    model: nn.Module,
    optimizer: Optional[torch.optim.Optimizer] = None,
    meta: Optional[Dict[str, Any]] = None,
) -> None:
    """Save ``model`` (and optionally ``optimizer``/``meta``) to ``path``.

    Parent directories are created if needed.
    """
    parent = os.path.dirname(os.path.abspath(path))
    os.makedirs(parent, exist_ok=True)
    payload: Dict[str, Any] = {"model_state": model.state_dict()}
    if optimizer is not None:
        payload["optimizer_state"] = optimizer.state_dict()
    if meta is not None:
        payload["meta"] = meta
    torch.save(payload, path)

train_teacher

train_teacher(teacher: Module, students: Sequence[Module], *, sample_batch: SampleBatch, build_prototype: BuildPrototype, forward: Forward, steps: int, lr: float = 0.001, proto_noise_std: float = 0.0, grad_clip: Optional[float] = 1.0, cosine: bool = True, extra_params: Optional[Sequence[Parameter]] = None, log_every: int = 0, log_fn: Callable[[str], None] = print) -> TeacherTrainResult

Train teacher to generate weights for a population of students.

Parameters:

Name	Type	Description	Default
teacher	Module	the hypernetwork teacher (must expose pi_scale_mean() if you want the pi diagnostic recorded; it is ignored when it returns None).	required
students	Sequence[Module]	frozen students sampled uniformly each step.	required
sample_batch	SampleBatch	returns a fresh training batch.	required
build_prototype	BuildPrototype	maps (student, batch) to a prototype tensor.	required
forward	Forward	maps (student, batch, generated_weights) to (logits, target) for the cross-entropy loss.	required
steps	int	number of optimisation steps.	required
lr	float	Adam learning rate.	0.001
proto_noise_std	float	std of optional Gaussian noise added to the prototype (a regulariser; the paper's matched regime uses 0).	0.0
grad_clip	Optional[float]	gradient-norm clip value (None to disable).	1.0
cosine	bool	use cosine LR annealing over steps (default True).	True
extra_params	Optional[Sequence[Parameter]]	additional parameters to optimise jointly with the teacher (e.g. a :class:LearnablePrototypeEncoder's parameters).	None
log_every	int	print a progress line every log_every steps (0 = silent).	0
log_fn	Callable[[str], None]	sink for log lines (default print).	print

Returns:

Name	Type	Description
A	TeacherTrainResult	class:TeacherTrainResult with per-step losses and pi-scale values.

Source code in polyweave/training/loop.py

def train_teacher(
    teacher: nn.Module,
    students: Sequence[nn.Module],
    *,
    sample_batch: SampleBatch,
    build_prototype: BuildPrototype,
    forward: Forward,
    steps: int,
    lr: float = 1e-3,
    proto_noise_std: float = 0.0,
    grad_clip: Optional[float] = 1.0,
    cosine: bool = True,
    extra_params: Optional[Sequence[nn.Parameter]] = None,
    log_every: int = 0,
    log_fn: Callable[[str], None] = print,
) -> TeacherTrainResult:
    """Train ``teacher`` to generate weights for a population of ``students``.

    Args:
        teacher: the hypernetwork teacher (must expose ``pi_scale_mean()`` if you
            want the pi diagnostic recorded; it is ignored when it returns None).
        students: frozen students sampled uniformly each step.
        sample_batch: returns a fresh training batch.
        build_prototype: maps ``(student, batch)`` to a prototype tensor.
        forward: maps ``(student, batch, generated_weights)`` to
            ``(logits, target)`` for the cross-entropy loss.
        steps: number of optimisation steps.
        lr: Adam learning rate.
        proto_noise_std: std of optional Gaussian noise added to the prototype
            (a regulariser; the paper's matched regime uses 0).
        grad_clip: gradient-norm clip value (None to disable).
        cosine: use cosine LR annealing over ``steps`` (default True).
        extra_params: additional parameters to optimise jointly with the teacher
            (e.g. a :class:`LearnablePrototypeEncoder`'s parameters).
        log_every: print a progress line every ``log_every`` steps (0 = silent).
        log_fn: sink for log lines (default ``print``).

    Returns:
        A :class:`TeacherTrainResult` with per-step losses and pi-scale values.
    """
    for s in students:
        s.eval()
        for p in s.parameters():
            p.requires_grad_(False)

    params = list(teacher.parameters())
    if extra_params is not None:
        params += list(extra_params)
    opt = torch.optim.Adam(params, lr=lr)
    sched = (
        torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=steps) if cosine else None
    )

    pi_fn = getattr(teacher, "pi_scale_mean", None)
    exp_fn = getattr(teacher, "exponent_abs_mean", None)  # metric A (weights only)
    share_fn = getattr(teacher, "pi_share", None)         # metric B (needs proto)
    result = TeacherTrainResult()

    teacher.train()
    for step in range(1, steps + 1):
        batch = sample_batch()
        student = random.choice(list(students))
        proto = build_prototype(student, batch)
        if proto_noise_std > 0:
            proto = proto + proto_noise_std * torch.randn_like(proto)

        gen = teacher(proto)
        logits, target = forward(student, batch, gen)
        loss = F.cross_entropy(logits, target)

        opt.zero_grad()
        loss.backward()
        if grad_clip is not None:
            torch.nn.utils.clip_grad_norm_(params, grad_clip)
        opt.step()
        if sched is not None:
            sched.step()

        result.losses.append(loss.item())
        if pi_fn is not None:
            pv = pi_fn()
            if pv is not None:
                result.pi_scales.append(pv)
        if exp_fn is not None:
            ev = exp_fn()
            if ev is not None:
                result.exponent_abs_means.append(ev)
        if share_fn is not None:
            sv = share_fn(proto)
            if sv is not None:
                result.pi_shares.append(sv)

        if log_every and step % log_every == 0:
            acc = (logits.argmax(1) == target).float().mean().item()
            extra = ""
            if result.pi_scales:
                extra += f"  pi_scale={result.pi_scales[-1]:.5f}"
            if result.exponent_abs_means:
                extra += f"  A|exp|={result.exponent_abs_means[-1]:.5f}"
            if result.pi_shares:
                extra += f"  B_share={result.pi_shares[-1]:.4f}"
            log_fn(f"  step {step:5d}/{steps}: loss={loss.item():.4f}  acc={acc:.3f}{extra}")

    return result