"""PyTorch backend: wraps the CIF obs operator as a ``torch.autograd.Function``.
Training modes
--------------
Mode 2 (online)
CIF called live as teacher; ``adjoint_consistency=False``.
Mode 3 (adjoint)
Adds a VJP-agreement loss between the surrogate Jacobian and the CIF
adjoint; ``adjoint_consistency=True``. Uses ``create_graph=True`` so
the adjoint-consistency loss is itself differentiable w.r.t. model params.
ONNX export
-----------
Uses ``torch.onnx.export`` with opset 17. The batch axis is declared dynamic
so the resulting file accepts any batch size at inference time.
"""
from __future__ import annotations
import os
import numpy as np
from typing import Any, Callable, List, Tuple
from .base import SurrogateBackend
def _sanitize_omp_num_threads() -> None:
"""Ensure OMP_NUM_THREADS is a plain integer before PyTorch is imported.
On some HPC schedulers (e.g. SLURM) the variable is set to a range such
as ``"1:2"``, which is invalid for OpenMP and causes PyTorch's thread-pool
initialisation to abort with an out-of-heap-memory error. We parse the
value and keep only the first integer, defaulting to 1 if parsing fails.
"""
raw = os.environ.get("OMP_NUM_THREADS", "")
if not raw:
return
try:
int(raw) # already valid — nothing to do
except ValueError:
# Take the first token that looks like an integer (e.g. "1" from "1:2")
first = raw.split(":")[0].strip()
os.environ["OMP_NUM_THREADS"] = first if first.isdigit() else "1"
[docs]
class TorchBackend(SurrogateBackend):
"""PyTorch surrogate backend."""
# ------------------------------------------------------------------
[docs]
def build_default_model(
self,
n_state: int,
n_obs: int,
hidden_size: int,
n_hidden_layers: int,
) -> Any:
"""Return a ``torch.nn.Sequential`` MLP: n_state -> hidden -> n_obs."""
_sanitize_omp_num_threads()
try:
import torch.nn as nn
except ImportError as exc:
raise ImportError(
"PyTorch is required for the torch backend: pip install torch"
) from exc
layers = [nn.Linear(n_state, hidden_size), nn.ReLU()]
for _ in range(n_hidden_layers - 1):
layers += [nn.Linear(hidden_size, hidden_size), nn.ReLU()]
layers.append(nn.Linear(hidden_size, n_obs))
return nn.Sequential(*layers)
# ------------------------------------------------------------------
[docs]
def wrap_operator(self, run_mode: str = "fwd",
reload_results: bool = False) -> Callable:
"""Return a ``torch.autograd.Function`` wrapping the CIF operator.
Args:
run_mode: ``"fwd"`` uses ``obs_operator.forward``; ``"tl"`` uses
``obs_operator.tangent_linear``. The backward pass always
calls ``obs_operator.adjoint``.
reload_results: Forwarded to every operator call.
"""
_sanitize_omp_num_threads()
try:
import torch
except ImportError as exc:
raise ImportError(
"PyTorch is required for the torch backend: pip install torch"
) from exc
obs_operator = self.obs_operator
fwd_call = (obs_operator.tangent_linear
if run_mode == "tl" else obs_operator.forward)
class _CIFFunction(torch.autograd.Function):
@staticmethod
def forward(ctx: Any, x: "torch.Tensor") -> "torch.Tensor":
x_np = x.detach().cpu().numpy().astype(np.float64)
y_np = fwd_call(x_np, reload_results=reload_results)
ctx._x_np = x_np
return torch.as_tensor(y_np, dtype=x.dtype, device=x.device)
@staticmethod
def backward(
ctx: Any, grad_output: "torch.Tensor"
) -> "torch.Tensor":
dy_np = grad_output.detach().cpu().numpy().astype(np.float64)
dx_np = obs_operator.adjoint(ctx._x_np, dy_np,
reload_results=reload_results)
return torch.as_tensor(
dx_np, dtype=grad_output.dtype, device=grad_output.device
)
def H(x: "torch.Tensor") -> "torch.Tensor":
return _CIFFunction.apply(x)
return H
# ------------------------------------------------------------------
[docs]
def train(
self,
model: Any,
sample_fn: Callable[[int], Any],
n_steps: int = 5000,
*,
run_mode: str = "fwd",
reload_results: bool = False,
adjoint_consistency: bool = False,
lambda_adj: float = 0.1,
) -> Tuple[Any, List[float]]:
"""Train *model* against the live CIF operator.
Args:
model: A ``torch.nn.Module``.
sample_fn: ``callable(batch_size) -> Tensor | ndarray`` returning
a single state vector of shape ``(N_state,)``.
n_steps: Number of gradient steps.
run_mode: Forwarded to :meth:`wrap_operator`.
reload_results: Forwarded to :meth:`wrap_operator`.
adjoint_consistency: Add VJP-agreement loss term (Mode 3).
lambda_adj: Weight of the adjoint consistency loss.
Returns:
``(trained_model, losses)``
"""
try:
import torch
import torch.nn.functional as F
except ImportError as exc:
raise ImportError(
"PyTorch is required for the torch backend: pip install torch"
) from exc
obs_operator = self.obs_operator
H = self.wrap_operator(run_mode=run_mode, reload_results=reload_results)
optimizer = torch.optim.Adam(model.parameters())
losses: List[float] = []
model.train()
for _ in range(n_steps):
raw = sample_fn(1)
if not isinstance(raw, torch.Tensor):
x = torch.as_tensor(raw, dtype=torch.float32)
else:
x = raw.float()
# Mode 2: CIF teacher — no gradient through CIF
with torch.no_grad():
y_target = H(x)
y_pred = model(x)
loss = F.mse_loss(y_pred, y_target)
if adjoint_consistency:
# Mode 3: match surrogate VJP to CIF adjoint
dy = torch.randn_like(y_target)
dx_cif = torch.as_tensor(
obs_operator.adjoint(
x.detach().cpu().numpy().astype(np.float64),
dy.detach().cpu().numpy().astype(np.float64),
),
dtype=x.dtype,
)
x_req = x.detach().requires_grad_(True)
y_surr = model(x_req)
(dx_surr,) = torch.autograd.grad(
y_surr, x_req, grad_outputs=dy, create_graph=True
)
loss = loss + lambda_adj * F.mse_loss(dx_surr, dx_cif)
optimizer.zero_grad()
loss.backward()
optimizer.step()
losses.append(float(loss.item()))
return model, losses
# ------------------------------------------------------------------
[docs]
def export_onnx(
self, path: str, model: Any, input_shape: Tuple[int, ...]
) -> None:
"""Export *model* to ONNX via ``torch.onnx.export`` (opset 17).
The batch dimension (axis 0) is made dynamic so the ONNX file accepts
any batch size at inference time.
Args:
path: Destination ``.onnx`` file path.
model: A trained ``torch.nn.Module``.
input_shape: Shape of a single input including batch dim,
e.g. ``(1, N_state)``.
"""
try:
import torch
except ImportError as exc:
raise ImportError(
"PyTorch is required for ONNX export: pip install torch"
) from exc
model.eval()
dummy = torch.zeros(*input_shape)
torch.onnx.export(
model,
dummy,
path,
input_names=["input"],
output_names=["output"],
dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
opset_version=14,
)
