# Example for surrogate model with a custom architecture using `pytorch` This example shows how to define a `pytorch` model architecture 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`](./../Basics/campaign.md) for a basic example. ## Necessary imports ```python import numpy as np import torch from torch import Tensor, nn ``` ```python from baybe.campaign import Campaign from baybe.objectives import SingleTargetObjective from baybe.parameters import ( CategoricalParameter, NumericalDiscreteParameter, SubstanceParameter, ) from baybe.recommenders import ( BotorchRecommender, FPSRecommender, TwoPhaseMetaRecommender, ) from baybe.searchspace import SearchSpace from baybe.surrogates import register_custom_architecture from baybe.targets import NumericalTarget from baybe.utils.dataframe import add_fake_results ``` ```python torch.set_default_dtype(torch.float64) ``` ## Architecture definition Note that the following is an example `PyTorch` Neural Network. Details of the setup is not the focus of BayBE but can be found in `Pytorch` guides. Model Configuration ```python INPUT_DIM = 10 OUTPUT_DIM = 1 DROPOUT = 0.5 NUM_NEURONS = [128, 32, 8] ``` Model training hyperparameters ```python HYPERPARAMS = { "epochs": 10, "lr": 1e-3, "criterion": nn.MSELoss, "optimizer": torch.optim.Adam, } ``` MC Parameters ```python MC = 100 ``` Helper functions ```python def _create_linear_block(in_features: int, out_features: int) -> list: """Create a linear block with dropout and relu activation.""" return [nn.Linear(in_features, out_features), nn.Dropout(p=DROPOUT), nn.ReLU()] ``` ```python def _create_hidden_layers(num_neurons: list[int]) -> list: """Create all hidden layers comprised of linear blocks.""" layers = [] for in_features, out_features in zip(num_neurons, num_neurons[1:]): layers.extend(_create_linear_block(in_features, out_features)) return layers ``` Model Architecture ```python class NeuralNetDropout(nn.Module): """Pytorch implementation of Neural Network with Dropout.""" def __init__(self): super().__init__() layers = [ # Initial linear block with input *(_create_linear_block(INPUT_DIM, NUM_NEURONS[0])), # All hidden layers *(_create_hidden_layers(NUM_NEURONS)), # Last linear output nn.Linear(NUM_NEURONS[-1], OUTPUT_DIM), ] # Sequential with layers (Feed Forward) self.model = nn.Sequential(*layers) def forward(self, data: Tensor) -> Tensor: """Forward method for NN.""" return self.model(data) ``` ## Surrogate Definition with BayBE Registration The class must include `_fit` and `_posterior` functions with the correct signatures Registration ```python @register_custom_architecture( joint_posterior_attr=False, constant_target_catching=False, batchify_posterior=True ) class NeuralNetDropoutSurrogate: """Surrogate that extracts posterior using monte carlo dropout simulations.""" def __init__(self): self.model: nn.Module | None = None def _posterior(self, candidates: Tensor) -> tuple[Tensor, Tensor]: """See :class:`baybe.surrogates.Surrogate`.""" self.model = self.model.train() # keep dropout # Run mc experiments through the NN with dropout predictions = torch.cat( [self.model(candidates).unsqueeze(dim=0) for _ in range(MC)] ) # Compute posterior mean and variance mean = predictions.mean(dim=0) var = predictions.var(dim=0) return mean, var def _fit(self, searchspace: SearchSpace, train_x: Tensor, train_y: Tensor) -> None: """See :class:`baybe.surrogates.Surrogate`.""" # Initialize Model self.model = NeuralNetDropout() # Training hyperparameters opt = HYPERPARAMS["optimizer"](self.model.parameters(), lr=HYPERPARAMS["lr"]) criterion = HYPERPARAMS["criterion"]() # Training loop for _ in range(HYPERPARAMS["epochs"]): opt.zero_grad() preds = self.model(train_x) loss = criterion(preds, train_y) loss.backward() opt.step() ``` ## Experiment Setup ```python parameters = [ CategoricalParameter( name="Granularity", values=["coarse", "medium", "fine"], encoding="OHE", ), NumericalDiscreteParameter( name="Pressure[bar]", values=[1, 5, 10], tolerance=0.2, ), NumericalDiscreteParameter( name="Temperature[degree_C]", values=np.linspace(100, 200, 10), ), SubstanceParameter( name="Solvent", data={ "Solvent A": "COC", "Solvent B": "CCC", "Solvent C": "O", "Solvent D": "CS(=O)C", }, encoding="MORDRED", ), ] ``` ## Run DOE iterations with custom surrogate Create campaign ```python campaign = Campaign( searchspace=SearchSpace.from_product(parameters=parameters, constraints=None), objective=SingleTargetObjective(target=NumericalTarget(name="Yield", mode="MAX")), recommender=TwoPhaseMetaRecommender( recommender=BotorchRecommender(surrogate_model=NeuralNetDropoutSurrogate()), initial_recommender=FPSRecommender(), ), ) ``` ________________________________________________________________________________ [Memory] Calling baybe.utils.chemistry._smiles_to_mordred_features... _smiles_to_mordred_features('CS(=O)C') _______________________________________smiles_to_mordred_features - 0.1s, 0.0min ```python # Let's do a first round of recommendation recommendation = campaign.recommend(batch_size=2) ``` ```python print("Recommendation from campaign:") print(recommendation) ``` Recommendation from campaign: Granularity Pressure[bar] Temperature[degree_C] Solvent 2 coarse 1.0 100.0 Solvent C 239 medium 10.0 200.0 Solvent D Add some fake results ```python add_fake_results(recommendation, campaign.targets) campaign.add_measurements(recommendation) ``` ```python # Do another round of recommendations recommendation = campaign.recommend(batch_size=2) ``` Print second round of recommendations ```python print("Recommendation from campaign:") print(recommendation) ``` Recommendation from campaign: Granularity Pressure[bar] Temperature[degree_C] Solvent index 140 medium 1.0 155.555556 Solvent A 196 medium 5.0 200.000000 Solvent A ```python print() ``` ## Serialization Serialization of custom models is not supported ```python try: campaign.to_json() except RuntimeError as e: print(f"Serialization Error Message: {e}") ``` Serialization Error Message: Serializing objects of type 'CustomArchitectureSurrogate' is not supported.