From 4cb522bbe8eab566525ea173b41edccaab180386 Mon Sep 17 00:00:00 2001 From: simonsung06 <67093926+simonsung06@users.noreply.github.com> Date: Tue, 7 Nov 2023 18:07:46 +0800 Subject: [PATCH] add output_scaler to mlp (#305) Co-authored-by: Simon --- bofire/data_models/surrogates/mlp.py | 21 +++++++++++++++++++- bofire/surrogates/mlp.py | 18 ++++++++++++++--- tests/bofire/data_models/specs/surrogates.py | 1 + tests/bofire/surrogates/test_mlp.py | 18 +++++++++++++++-- 4 files changed, 52 insertions(+), 6 deletions(-) diff --git a/bofire/data_models/surrogates/mlp.py b/bofire/data_models/surrogates/mlp.py index 57bd61bc1..72336334e 100644 --- a/bofire/data_models/surrogates/mlp.py +++ b/bofire/data_models/surrogates/mlp.py @@ -1,6 +1,6 @@ from typing import Annotated, Literal, Sequence -from pydantic import Field +from pydantic import Field, validator from bofire.data_models.surrogates.botorch import BotorchSurrogate from bofire.data_models.surrogates.scaler import ScalerEnum @@ -20,3 +20,22 @@ class MLPEnsemble(BotorchSurrogate, TrainableSurrogate): subsample_fraction: Annotated[float, Field(gt=0.0)] = 1.0 shuffle: bool = True scaler: ScalerEnum = ScalerEnum.NORMALIZE + output_scaler: ScalerEnum = ScalerEnum.STANDARDIZE + + @validator("output_scaler") + def validate_output_scaler(cls, output_scaler): + """validates that output_scaler is a valid type + + Args: + output_scaler (ScalerEnum): Scaler used to transform the output + + Raises: + ValueError: when ScalerEnum.NORMALIZE is used + + Returns: + ScalerEnum: Scaler used to transform the output + """ + if output_scaler == ScalerEnum.NORMALIZE: + raise ValueError("Normalize is not supported as an output transform.") + + return output_scaler diff --git a/bofire/surrogates/mlp.py b/bofire/surrogates/mlp.py index cb4cdc18d..28b3eb87f 100644 --- a/bofire/surrogates/mlp.py +++ b/bofire/surrogates/mlp.py @@ -5,11 +5,13 @@ import torch import torch.nn as nn from botorch.models.ensemble import EnsembleModel +from botorch.models.transforms.outcome import OutcomeTransform, Standardize from torch import Tensor from torch.utils.data import DataLoader, Dataset from bofire.data_models.enum import OutputFilteringEnum from bofire.data_models.surrogates.api import MLPEnsemble as DataModel +from bofire.data_models.surrogates.scaler import ScalerEnum from bofire.surrogates.botorch import BotorchSurrogate from bofire.surrogates.single_task_gp import get_scaler from bofire.surrogates.trainable import TrainableSurrogate @@ -74,7 +76,9 @@ def forward(self, x): class _MLPEnsemble(EnsembleModel): - def __init__(self, mlps: Sequence[MLP]): + def __init__( + self, mlps: Sequence[MLP], output_scaler: Optional[OutcomeTransform] = None + ): super().__init__() if len(mlps) == 0: raise ValueError("List of mlps is empty.") @@ -84,6 +88,8 @@ def __init__(self, mlps: Sequence[MLP]): assert mlp.layers[0].in_features == num_in_features assert mlp.layers[-1].out_features == num_out_features self.mlps = mlps + if output_scaler is not None: + self.outcome_transform = output_scaler # put all models in eval mode for mlp in self.mlps: mlp.eval() @@ -170,6 +176,7 @@ def __init__(self, data_model: DataModel, **kwargs): self.subsample_fraction = data_model.subsample_fraction self.shuffle = data_model.shuffle self.scaler = data_model.scaler + self.output_scaler = data_model.output_scaler super().__init__(data_model, **kwargs) _output_filtering: OutputFilteringEnum = OutputFilteringEnum.ALL @@ -179,6 +186,11 @@ def _fit(self, X: pd.DataFrame, Y: pd.DataFrame): scaler = get_scaler(self.inputs, self.input_preprocessing_specs, self.scaler, X) transformed_X = self.inputs.transform(X, self.input_preprocessing_specs) + if self.output_scaler == ScalerEnum.STANDARDIZE: + output_scaler = Standardize(m=Y.shape[-1]) + else: + output_scaler = None + mlps = [] subsample_size = round(self.subsample_fraction * X.shape[0]) for _ in range(self.n_estimators): @@ -189,7 +201,7 @@ def _fit(self, X: pd.DataFrame, Y: pd.DataFrame): dataset = RegressionDataSet( X=scaler.transform(tX) if scaler is not None else tX, - y=ty, + y=output_scaler(ty)[0] if output_scaler is not None else ty, ) mlp = MLP( input_size=transformed_X.shape[1], @@ -208,6 +220,6 @@ def _fit(self, X: pd.DataFrame, Y: pd.DataFrame): weight_decay=self.weight_decay, ) mlps.append(mlp) - self.model = _MLPEnsemble(mlps=mlps) + self.model = _MLPEnsemble(mlps, output_scaler=output_scaler) if scaler is not None: self.model.input_transform = scaler diff --git a/tests/bofire/data_models/specs/surrogates.py b/tests/bofire/data_models/specs/surrogates.py index fca8afa6b..ccbffb678 100644 --- a/tests/bofire/data_models/specs/surrogates.py +++ b/tests/bofire/data_models/specs/surrogates.py @@ -176,6 +176,7 @@ "subsample_fraction": 1.0, "shuffle": True, "scaler": ScalerEnum.NORMALIZE, + "output_scaler": ScalerEnum.STANDARDIZE, "input_preprocessing_specs": {}, "dump": None, "hyperconfig": None, diff --git a/tests/bofire/surrogates/test_mlp.py b/tests/bofire/surrogates/test_mlp.py index 26ff87744..8acde7c63 100644 --- a/tests/bofire/surrogates/test_mlp.py +++ b/tests/bofire/surrogates/test_mlp.py @@ -2,6 +2,7 @@ import torch import torch.nn as nn from botorch.models.transforms.input import InputStandardize, Normalize +from botorch.models.transforms.outcome import Standardize from pandas.testing import assert_frame_equal import bofire.surrogates.api as surrogates @@ -163,9 +164,14 @@ def test_mlp_ensemble_forward(): @pytest.mark.parametrize( - "scaler", [ScalerEnum.NORMALIZE, ScalerEnum.STANDARDIZE, ScalerEnum.IDENTITY] + "scaler, output_scaler", + [ + [ScalerEnum.NORMALIZE, ScalerEnum.IDENTITY], + [ScalerEnum.STANDARDIZE, ScalerEnum.STANDARDIZE], + [ScalerEnum.IDENTITY, ScalerEnum.STANDARDIZE], + ], ) -def test_mlp_ensemble_fit(scaler): +def test_mlp_ensemble_fit(scaler, output_scaler): bench = Himmelblau() samples = bench.domain.inputs.sample(10) experiments = bench.f(samples, return_complete=True) @@ -175,10 +181,12 @@ def test_mlp_ensemble_fit(scaler): n_estimators=2, n_epochs=5, scaler=scaler, + output_scaler=output_scaler, ) surrogate = surrogates.map(ens) surrogate.fit(experiments=experiments) + if scaler == ScalerEnum.NORMALIZE: assert isinstance(surrogate.model.input_transform, Normalize) elif scaler == ScalerEnum.STANDARDIZE: @@ -186,6 +194,12 @@ def test_mlp_ensemble_fit(scaler): else: with pytest.raises(AttributeError): assert surrogate.model.input_transform is None + + if output_scaler == ScalerEnum.STANDARDIZE: + assert isinstance(surrogate.model.outcome_transform, Standardize) + elif output_scaler == ScalerEnum.IDENTITY: + assert not hasattr(surrogate.model, "outcome_transform") + preds = surrogate.predict(experiments) dump = surrogate.dumps() surrogate2 = surrogates.map(ens)