polystep 0.7.0__py3-none-any.whl

polystep/api.py

@@ -0,0 +1,307 @@
+"""High-level training API for polystep.
+
+Provides a single-call ``train()`` function that wraps PolyStepOptimizer
+with an epoch-based training loop, callback system, and diagnostics.
+Users can train any ``nn.Module`` without manually managing closures,
+epoch iteration, or batch handling.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, List, Optional
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+
+from .cost_nn import NNCostEvaluator
+from .optimizer import PolyStepOptimizer
+
+
+class TrainCallback:
+    """Base class for training callbacks.
+
+    Subclass and override ``on_step_end`` and/or ``on_epoch_end`` to
+    hook into the training loop. Return ``True`` from ``on_step_end``
+    to signal early stopping.
+
+    Example::
+
+        class LossThresholdCallback(TrainCallback):
+            def __init__(self, threshold: float = 0.01):
+                self.threshold = threshold
+
+            def on_step_end(self, metrics: dict) -> bool:
+                if metrics['loss'] < self.threshold:
+                    print(f"Target loss reached at step {metrics['step']}")
+                    return True  # stop training
+                return False
+    """
+
+    def on_step_end(self, metrics: dict) -> bool:
+        """Called after each optimizer step.
+
+        Args:
+            metrics: Dict with keys 'step', 'epoch', 'loss', 'ot_cost',
+                'displacement', 'velocity_mag', 'converged'.
+
+        Returns:
+            True to stop training, False to continue.
+        """
+        return False
+
+    def on_epoch_end(self, metrics: dict) -> None:
+        """Called after each epoch completes.
+
+        Args:
+            metrics: Dict with keys 'epoch', 'avg_loss'.
+        """
+        pass
+
+
+@dataclass
+class TrainConfig:
+    """Configuration for the ``train()`` function.
+
+    Training loop parameters only. Optimizer-specific parameters (epsilon,
+    step_radius, subspace, block_strategy, etc.) are set on
+    ``PolyStepOptimizer``.
+
+    Args:
+        epochs: Number of training epochs. Must be > 0.
+        batch_size: Convenience/documentation field. ``train()`` uses
+            the provided DataLoader directly and does NOT build one
+            from this value.
+        log_every: Step interval for built-in logging. Must be > 0.
+        callbacks: List of ``TrainCallback`` instances. ``None`` is
+            normalized to an empty list.
+    """
+
+    epochs: int = 10
+    batch_size: int = 32
+    log_every: int = 10
+    callbacks: Optional[List[TrainCallback]] = None
+
+    def __post_init__(self):
+        if self.epochs <= 0:
+            raise ValueError(f"epochs must be > 0, got {self.epochs}")
+        if self.batch_size <= 0:
+            raise ValueError(f"batch_size must be > 0, got {self.batch_size}")
+        if self.log_every <= 0:
+            raise ValueError(f"log_every must be > 0, got {self.log_every}")
+        if self.callbacks is None:
+            self.callbacks = []
+
+
+def train(
+    model: nn.Module,
+    dataloader: DataLoader,
+    loss_fn: Callable,
+    optimizer: PolyStepOptimizer,
+    config: TrainConfig,
+) -> nn.Module:
+    """Train a model using the Sinkhorn Step optimizer.
+
+    Builds the OT closure internally from the model, loss function, and
+    each mini-batch. Users never handle closures directly. For each batch,
+    the function constructs a closure that evaluates the model at multiple
+    candidate parameter configurations via ``NNCostEvaluator``, then calls
+    ``optimizer.step(closure)`` to update the model weights.
+
+    Example::
+
+        from polystep import PolyStepOptimizer, train, TrainConfig, LoggingCallback
+
+        optimizer = PolyStepOptimizer(model, epsilon=0.1)
+        config = TrainConfig(epochs=5, callbacks=[LoggingCallback()])
+        trained_model = train(model, dataloader, loss_fn, optimizer, config)
+
+    See Also:
+        ``PolyStepOptimizer`` for optimizer configuration (epsilon,
+        step_radius, subspace, block_strategy, etc.).
+
+    Args:
+        model: The ``nn.Module`` to train. Updated in-place.
+        dataloader: PyTorch DataLoader yielding ``(inputs, targets)`` batches.
+        loss_fn: Loss function with signature ``loss_fn(output, targets) -> scalar``.
+        optimizer: A ``PolyStepOptimizer`` instance (already configured).
+        config: ``TrainConfig`` with loop hyperparameters and callbacks.
+
+    Returns:
+        The same model object (mutated in-place by the optimizer).
+    """
+    evaluator = NNCostEvaluator(
+        model,
+        loss_fn=loss_fn,
+        compile_vmap=getattr(optimizer, "_compile_evaluator", False),
+    )
+    callbacks = list(config.callbacks)
+    global_step = 0
+    stop = False
+
+    # Detect model device for automatic batch transfer
+    try:
+        device = next(model.parameters()).device
+    except StopIteration:
+        raise ValueError("Model has no trainable parameters")
+
+    for epoch in range(config.epochs):
+        epoch_loss_sum = 0.0
+        epoch_loss_count = 0
+
+        for batch in dataloader:
+            inputs, targets = batch
+            inputs, targets = inputs.to(device), targets.to(device)
+
+            # Micro-batch: subsample for cost evaluation if configured.
+            cost_bs = getattr(optimizer, "cost_batch_size", None)
+            if cost_bs is not None and cost_bs < inputs.shape[0]:
+                gen = getattr(optimizer, "_generator", None)
+                if gen is not None and gen.device.type == inputs.device.type:
+                    idx = torch.randperm(inputs.shape[0], device=inputs.device, generator=gen)[:cost_bs]
+                else:
+                    idx = torch.randperm(inputs.shape[0], device=inputs.device)[:cost_bs]
+                cost_inputs = inputs[idx]
+                cost_targets = targets[idx]
+            else:
+                cost_inputs = inputs
+                cost_targets = targets
+
+            # Build closure for this batch using default-argument capture
+            def closure(batched_params, _in=cost_inputs, _tgt=cost_targets):
+                return evaluator.evaluate(batched_params, _in, _tgt)
+
+            optimizer.step(closure)
+
+            # Compute training loss separately (not the OT cost)
+            with torch.no_grad():
+                output = model(inputs)
+                train_loss = loss_fn(output, targets).item()
+
+            epoch_loss_sum += train_loss
+            epoch_loss_count += 1
+
+            # Build metrics dict
+            state = optimizer.state
+            metrics = {
+                "step": global_step,
+                "epoch": epoch,
+                "loss": train_loss,
+                "ot_cost": state.costs[-1] if state.costs else 0.0,
+                "displacement": (state.displacement_sqnorms[-1] if state.displacement_sqnorms else 0.0),
+                "velocity_mag": (torch.norm(state.velocity).item() if state.velocity is not None else 0.0),
+                "converged": (state.linear_convergence[-1] if state.linear_convergence else False),
+                "absorb_count": getattr(state, "absorb_count", 0),
+            }
+
+            # Invoke on_step_end callbacks
+            for cb in callbacks:
+                if cb.on_step_end(metrics):
+                    stop = True
+                    break
+
+            if stop:
+                break
+
+            global_step += 1
+
+        if stop:
+            break
+
+        # Epoch-level metrics
+        avg_loss = epoch_loss_sum / epoch_loss_count if epoch_loss_count > 0 else 0.0
+        epoch_metrics = {
+            "epoch": epoch,
+            "avg_loss": avg_loss,
+        }
+        for cb in callbacks:
+            cb.on_epoch_end(epoch_metrics)
+
+    return model
+
+
+# ---------------------------------------------------------------------------
+# Built-in callbacks
+# ---------------------------------------------------------------------------
+
+
+class LoggingCallback(TrainCallback):
+    """Prints step metrics at a configurable interval.
+
+    Args:
+        log_every: Print every *log_every* steps (default 10).
+    """
+
+    def __init__(self, log_every: int = 10):
+        self.log_every = log_every
+
+    def on_step_end(self, metrics: dict) -> bool:
+        if metrics["step"] % self.log_every == 0:
+            print(
+                f"[Step {metrics['step']}] "
+                f"loss={metrics['loss']:.4f} "
+                f"ot_cost={metrics['ot_cost']:.4f} "
+                f"disp={metrics['displacement']:.6f} "
+                f"converged={metrics['converged']}"
+            )
+        return False
+
+    def on_epoch_end(self, metrics: dict) -> None:
+        print(f"--- Epoch {metrics['epoch']} complete | avg_loss={metrics['avg_loss']:.4f} ---")
+
+
+class EarlyStoppingCallback(TrainCallback):
+    """Stops training when loss stagnates for *patience* steps.
+
+    Args:
+        patience: Number of steps without improvement before stopping.
+        min_delta: Minimum loss decrease to count as improvement.
+    """
+
+    def __init__(self, patience: int = 10, min_delta: float = 1e-4):
+        self.patience = patience
+        self.min_delta = min_delta
+        self.best_loss: float = float("inf")
+        self.counter: int = 0
+
+    def on_step_end(self, metrics: dict) -> bool:
+        loss = metrics["loss"]
+        if loss < self.best_loss - self.min_delta:
+            self.best_loss = loss
+            self.counter = 0
+        else:
+            self.counter += 1
+        if self.counter >= self.patience:
+            print(f"Early stopping at step {metrics['step']}")
+            return True
+        return False
+
+
+# ---------------------------------------------------------------------------
+# Diagnostics helper
+# ---------------------------------------------------------------------------
+
+
+def get_diagnostics(optimizer: PolyStepOptimizer) -> dict:
+    """Extract diagnostic summary from optimizer state.
+
+    Args:
+        optimizer: A ``PolyStepOptimizer`` instance (may have been stepped).
+
+    Returns:
+        Dict with keys: 'costs', 'displacement_sqnorms', 'convergence',
+        'velocity_magnitude', 'iteration_count', 'epsilon', 'radius_multiplier',
+        'absorb_count'.
+    """
+    state = optimizer.state
+    return {
+        "costs": list(state.costs),
+        "displacement_sqnorms": list(state.displacement_sqnorms),
+        "convergence": list(state.linear_convergence),
+        "velocity_magnitude": (torch.norm(state.velocity).item() if state.velocity is not None else None),
+        "iteration_count": state.iteration_count,
+        "epsilon": state.epsilon,
+        "radius_multiplier": state.radius_multiplier,
+        "absorb_count": getattr(state, "absorb_count", 0),
+    }

polystep/ask_tell.py

@@ -0,0 +1,164 @@
+"""PolyStep as a standard ask/tell optimizer.
+
+Exposes the Sinkhorn Step update behind the ``ask``/``tell`` interface used by
+evolution-strategy libraries (evosax, NeuroEvoBench, CMA-ES). ``ask`` returns
+candidate points (rotated polytope vertices around each particle) for the caller
+to evaluate; ``tell`` takes their fitness (lower is better), builds the cost
+matrix, and applies one softmax-weighted barycentric step. PolyStep is then a
+drop-in gradient-free optimizer comparable to any ES.
+"""
+
+from __future__ import annotations
+
+import math
+import warnings
+from typing import Callable, Optional, Union
+
+import torch
+
+from ._compiled import _barycentric_projection, _rotate_and_translate
+from .geometry import get_orthoplex_vertices, get_random_rotation_matrices
+from .solvers import SinkhornSolver, SoftmaxSolver
+
+__all__ = ["PolyStepES", "minimize"]
+
+
+class PolyStepES:
+    """Ask/tell wrapper around the Sinkhorn Step update.
+
+    Args:
+        dim: Dimensionality of a solution vector.
+        num_particles: Number of independent particles (each proposes a
+            polytope of candidates). Population size is ``num_particles * 2 * dim``.
+        epsilon: Entropic OT temperature for the softmax weighting.
+        step_radius: Geometric step size along the polytope directions.
+        solver: OT solver instance. Defaults to :class:`SoftmaxSolver`; pass a
+            :class:`SinkhornSolver` for the full entropic-OT plan.
+        scale_cost: Cost-matrix scaling passed to the solver ("mean", "max_cost",
+            a float divisor, or None).
+        x0: Initial position(s), shape ``(dim,)`` or ``(num_particles, dim)``.
+        seed: Seed for the rotation generator.
+        device, dtype: Tensor device and dtype.
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        num_particles: int = 1,
+        epsilon: float = 0.5,
+        step_radius: float = 0.5,
+        solver: Optional[Union[SoftmaxSolver, SinkhornSolver]] = None,
+        scale_cost: Optional[Union[str, float]] = "mean",
+        x0: Optional[torch.Tensor] = None,
+        seed: Optional[int] = None,
+        device: Union[str, torch.device] = "cpu",
+        dtype: torch.dtype = torch.float32,
+    ):
+        if dim < 1:
+            raise ValueError(f"dim must be >= 1, got {dim}.")
+        if num_particles < 1:
+            raise ValueError(f"num_particles must be >= 1, got {num_particles}.")
+        if not epsilon > 0:
+            raise ValueError(f"epsilon must be > 0, got {epsilon}.")
+        if not (math.isfinite(step_radius) and step_radius >= 0):
+            raise ValueError(f"step_radius must be finite and >= 0, got {step_radius}.")
+        self.dim = dim
+        self.num_particles = num_particles
+        self.epsilon = epsilon
+        self.step_radius = step_radius
+        self.scale_cost = scale_cost
+        self.device = torch.device(device)
+        self.dtype = dtype
+        self.solver = solver if solver is not None else SoftmaxSolver(epsilon=epsilon)
+        if isinstance(self.solver, SinkhornSolver) and num_particles == 1:
+            warnings.warn(
+                "SinkhornSolver with num_particles=1 yields a uniform transport plan "
+                "(the column marginal forces it), so steps ignore fitness. Use the "
+                "default SoftmaxSolver, or num_particles > 1.",
+                stacklevel=2,
+            )
+
+        self.generator = torch.Generator(device=self.device)
+        if seed is not None:
+            self.generator.manual_seed(seed)
+
+        self.vertices = get_orthoplex_vertices(dim, device=self.device, dtype=dtype)  # (2d, d)
+        self.num_vertices = self.vertices.shape[0]
+
+        if x0 is None:
+            X = torch.zeros(num_particles, dim, device=self.device, dtype=dtype)
+        else:
+            X = torch.as_tensor(x0, device=self.device, dtype=dtype).reshape(-1, dim)
+            if X.shape[0] == 1 and num_particles > 1:
+                X = X.expand(num_particles, dim).clone()
+        if X.shape[0] != num_particles:
+            raise ValueError(f"x0 has {X.shape[0]} particle rows but num_particles={num_particles}.")
+        self.X = X.clone()
+        self.a = torch.full((num_particles,), 1.0 / num_particles, device=self.device, dtype=dtype)
+
+        self._pending: Optional[torch.Tensor] = None
+        self.best_solution: Optional[torch.Tensor] = None
+        self.best_fitness: float = float("inf")
+
+    @property
+    def popsize(self) -> int:
+        return self.num_particles * self.num_vertices
+
+    @property
+    def mean(self) -> torch.Tensor:
+        """Mean particle position (the current solution estimate)."""
+        return self.X.mean(dim=0)
+
+    @torch.inference_mode()
+    def ask(self) -> torch.Tensor:
+        """Return candidate points of shape ``(popsize, dim)`` to evaluate."""
+        if self._pending is not None:
+            raise RuntimeError("ask() called twice before tell(); tell() the previous population first.")
+        rot = get_random_rotation_matrices(
+            self.num_particles, self.dim, device=self.device, dtype=self.dtype, generator=self.generator
+        )
+        X_vertices, _ = _rotate_and_translate(rot, self.vertices, self.X, self.step_radius)  # (P, V, d)
+        self._pending = X_vertices
+        return X_vertices.reshape(self.popsize, self.dim)
+
+    @torch.inference_mode()
+    def tell(self, fitness: torch.Tensor) -> None:
+        """Update particles from the fitness of the last ``ask`` (lower is better)."""
+        if self._pending is None:
+            raise RuntimeError("tell() called before ask()")
+        cost = torch.as_tensor(fitness, device=self.device, dtype=self.dtype).reshape(
+            self.num_particles, self.num_vertices
+        )
+        self.solver.epsilon = self.epsilon
+        # Source marginal defaults to uniform 1/P inside the solver, which is
+        # exactly ``self.a`` -- pass None so the solver skips the (host-syncing)
+        # user-marginal validation on this per-step path.
+        transport = self.solver.solve(cost, scale_cost=self.scale_cost).matrix
+        X_new = _barycentric_projection(transport, self.a, self._pending)
+        if torch.isfinite(X_new).all():
+            self.X = X_new
+
+        flat_cost = cost.reshape(-1)
+        fmin, idx = torch.min(flat_cost, dim=0)
+        if fmin.item() < self.best_fitness:
+            self.best_fitness = fmin.item()
+            self.best_solution = self._pending.reshape(self.popsize, self.dim)[idx].clone()
+        self._pending = None
+
+
+def minimize(
+    fn: Callable[[torch.Tensor], torch.Tensor],
+    dim: int,
+    steps: int = 200,
+    **kwargs,
+) -> PolyStepES:
+    """Minimize a batched black-box ``fn: (popsize, dim) -> (popsize,)``.
+
+    Runs ``steps`` ask/tell rounds and returns the optimizer, whose
+    ``best_solution`` / ``best_fitness`` / ``mean`` hold the result.
+    """
+    es = PolyStepES(dim, **kwargs)
+    for _ in range(steps):
+        candidates = es.ask()
+        es.tell(fn(candidates))
+    return es

polystep/benchmarks/__init__.py

@@ -0,0 +1,62 @@
+"""Benchmark utilities for polystep experiments.
+
+Provides shared model factories, data loaders, baselines, and evaluation
+helpers used by the experiment runners and examples.
+"""
+
+from .utils import (
+    SEEDS,
+    get_mnist_loaders,
+    get_cifar10_loaders,
+    get_sst2_loaders,
+    get_nmnist_loaders,
+    get_dvs_gesture_loaders,
+    MNISTNet,
+    CIFAR10Net,
+    create_model,
+    LIFNeuron,
+    SpikingNet,
+    evaluate_accuracy,
+    evaluate_snn_accuracy,
+    BenchmarkResult,
+    get_environment_info,
+    save_results_json,
+    format_results_table,
+    compute_summary_stats,
+)
+
+from .baselines import (
+    has_evotorch,
+    has_nevergrad,
+    train_cmaes,
+    train_nevergrad,
+    check_gradient_free_deps,
+    get_available_optimizers,
+)
+
+__all__ = [
+    "SEEDS",
+    "get_mnist_loaders",
+    "get_cifar10_loaders",
+    "get_sst2_loaders",
+    "get_nmnist_loaders",
+    "get_dvs_gesture_loaders",
+    "MNISTNet",
+    "CIFAR10Net",
+    "create_model",
+    "LIFNeuron",
+    "SpikingNet",
+    "evaluate_accuracy",
+    "evaluate_snn_accuracy",
+    "BenchmarkResult",
+    "get_environment_info",
+    "save_results_json",
+    "format_results_table",
+    "compute_summary_stats",
+    "has_evotorch",
+    "has_nevergrad",
+    "train_cmaes",
+    "train_nevergrad",
+    "check_gradient_free_deps",
+    "get_available_optimizers",
+]