Example for surrogate model with a custom pretrained model

This example shows how to pre-train a model and use it as a surrogate. Please note that the model is not designed to be useful but to demonstrate the workflow.

This example assumes some basic familiarity with using BayBE. We thus refer to campaign for a basic example.

Necessary imports

import numpy as np
import torch
from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType
from skl2onnx.operator_converters.linear_regressor import convert_sklearn_bayesian_ridge
from sklearn.linear_model import BayesianRidge

from baybe.campaign import Campaign
from baybe.objectives import SingleTargetObjective
from baybe.parameters import NumericalDiscreteParameter
from baybe.recommenders import (
    BotorchRecommender,
    FPSRecommender,
    TwoPhaseMetaRecommender,
)
from baybe.searchspace import SearchSpace
from baybe.surrogates import CustomONNXSurrogate
from baybe.targets import NumericalTarget
from baybe.utils.dataframe import add_fake_measurements, to_tensor

Experiment Setup

parameters = [
    NumericalDiscreteParameter(
        name="Pressure[bar]",
        values=[1, 5, 10],
        tolerance=0.2,
    ),
    NumericalDiscreteParameter(
        name="Temperature[degree_C]",
        values=np.linspace(100, 200, 10),
    ),
]

“Pre-training” stage

Note that this example trains with several helpers built-in to BayBE. This can be done independently (and elsewhere). The only requirement that BayBE needs is that the model is in an onnx format. And The format should return both the mean and standard deviation. This example is based on a BayesianRidge regressor from sklearn. Its native conversion to onnx is supported via skl2onnx. Please also note that this example does not give a useful model. Its purpose is to show the workflow for using pre-trained surrogates in BayBE.

searchspace = SearchSpace.from_product(parameters=parameters, constraints=None)
train_x = to_tensor(searchspace.discrete.comp_rep)
train_y = torch.rand(train_x.size(dim=0))  # train with a random y vector

Define model and fit

model = BayesianRidge()
model.fit(train_x, train_y)