saealib.TorchSurrogate#

class saealib.TorchSurrogate(model, optimizer_cls=None, optimizer_kwargs=None, loss_fn=None, epochs=100)[source]#

Bases: RegressionSurrogate

Surrogate adapter for PyTorch nn.Module models.

Wraps any torch.nn.Module that maps (n_samples, n_features) to (n_samples, n_obj). The model is trained from scratch on each fit call by resetting to the weights captured at construction time, so this adapter is safe to use with LocalSurrogateManager (which re-fits per candidate).

Parameters:

  • model (torch.nn.Module) – Neural network model. Its output shape must be (n_samples, n_obj).

  • optimizer_cls (type or None) – Optimizer class (e.g. torch.optim.Adam). Defaults to torch.optim.Adam.

  • optimizer_kwargs (dict or None) – Keyword arguments forwarded to the optimizer constructor. Defaults to {}.

  • loss_fn (callable or None) – Loss function. Defaults to torch.nn.MSELoss().

  • epochs (int) – Number of training epochs per fit call. Default: 100.

Raises:

ImportError – If PyTorch is not installed.

Methods

__init__

contract

Return the surrogate model contract.

fit

Fit the surrogate model.

post_fit

Post-fit lifecycle hook; override to inject custom processing.

predict

Predict using the surrogate model.

with_post_fit

Return a copy of this surrogate with fn appended to the post-fit hook.

Method Details

TorchSurrogate.__init__(model, optimizer_cls=None, optimizer_kwargs=None, loss_fn=None, epochs=100)[source]#
Parameters:

  • model (torch.nn.Module)

  • optimizer_cls (type | None)

  • optimizer_kwargs (dict | None)

  • loss_fn (Callable[..., Any] | None)

  • epochs (int)

Return type:

None

TorchSurrogate.contract()#

Return the surrogate model contract.

Return type:

ComponentContract

TorchSurrogate.fit(train_x, train_y)[source]#

Fit the surrogate model.

Resets model weights to their initial values before training.

Parameters:

  • train_x (np.ndarray) – Training input data. shape: (n_samples, n_features)

  • train_y (np.ndarray) – Training output data. shape: (n_samples, n_obj) or (n_samples,).

Return type:

None

TorchSurrogate.post_fit(train_x, train_y, ctx=None)#

Post-fit lifecycle hook; override to inject custom processing.

Parameters:

  • train_x (np.ndarray) – Training input data used for this fit. shape: (n_samples, n_features)

  • train_y (np.ndarray) – Training output data used for this fit. shape: (n_samples, n_obj)

  • ctx (OptimizationState or None, optional) – Current optimization context.

Return type:

None

TorchSurrogate.predict(test_x)[source]#

Predict using the surrogate model.

Parameters:

test_x (np.ndarray) – Input data. shape: (n_samples, n_features) or (n_features,)

Returns:

prediction.value shape: (n_samples, n_obj) prediction.std is None

Return type:

SurrogatePrediction

TorchSurrogate.with_post_fit(fn)#

Return a copy of this surrogate with fn appended to the post-fit hook.

Parameters:

fn (callable) – fn(train_x, train_y, ctx) -> None

Returns:

Shallow copy with the hook registered.

Return type:

Surrogate