Introduce ModelDict and MultiModelAcquisitionFunction (#1584)

Summary:
Pull Request resolved: #1584

Introduces a lightweight `ModelDict` container, which is simply a `ModuleDict[str, Model]`, and an abstract `MultiModelAcquisitionFunction` class that accepts a `ModelDict` rather than a `Model`.

The goal here is to help shape the MBM Surrogate refactor by having a concrete example of how the multiple surrogates would be consumed in BoTorch.

Reviewed By: lena-kashtelyan

Differential Revision: D41564744

fbshipit-source-id: d18632b94041d3529cc0b98560d7791e7d83d59d

Author: saitcakmak

botorch/acquisition/acquisition.py

@@ -14,7 +14,7 @@
 
 import torch
 from botorch.exceptions import BotorchWarning, UnsupportedError
-from botorch.models.model import Model
+from botorch.models.model import Model, ModelDict
 from botorch.posteriors.posterior import Posterior
 from botorch.sampling.base import MCSampler
 from botorch.sampling.get_sampler import get_sampler
@@ -168,3 +168,34 @@ def get_posterior_samples(self, posterior: Posterior) -> Tensor:
                 posterior=posterior, sample_shape=self._default_sample_shape
             )
         return self.sampler(posterior=posterior)
+
+
+class MultiModelAcquisitionFunction(AcquisitionFunction, ABC):
+    r"""Abstract base class for acquisition functions that require
+    multiple types of models.
+
+    The intended use case for these acquisition functions are those
+    where we have multiple models, each serving a distinct purpose.
+    As an example, we can have a "regression" model that predicts
+    one or more outcomes, and a "classification" model that predicts
+    the probabilty that a given parameterization is feasible. The
+    multi-model acquisition function can then weight the acquisition
+    value computed with the "regression" model with the feasibility
+    value predicted by the "classification" model to produce the
+    composite acquisition value.
+
+    This is currently only a placeholder to help with some development
+    in Ax. We plan to add some acquisition functions utilizing multiple
+    models in the future.
+
+    :meta private:
+    """
+
+    def __init__(self, model_dict: ModelDict) -> None:
+        r"""Constructor for the MultiModelAcquisitionFunction base class.
+
+        Args:
+            model_dict: A ModelDict mapping labels to models.
+        """
+        super(AcquisitionFunction, self).__init__()
+        self.model_dict: ModelDict = model_dict

botorch/models/model.py

@@ -32,13 +32,14 @@
 import numpy as np
 import torch
 from botorch import settings
+from botorch.exceptions.errors import InputDataError
 from botorch.models.utils.assorted import fantasize as fantasize_flag
 from botorch.posteriors import Posterior, PosteriorList
 from botorch.sampling.base import MCSampler
 from botorch.utils.datasets import BotorchDataset
 from botorch.utils.transforms import is_fully_bayesian
 from torch import Tensor
-from torch.nn import Module, ModuleList
+from torch.nn import Module, ModuleDict, ModuleList
 
 if TYPE_CHECKING:
     from botorch.acquisition.objective import PosteriorTransform  # pragma: no cover
@@ -514,3 +515,20 @@ def load_state_dict(
                 }
                 m.load_state_dict(filtered_dict)
         super().load_state_dict(state_dict=state_dict, strict=strict)
+
+
+class ModelDict(ModuleDict):
+    r"""A lightweight container mapping model names to models."""
+
+    def __init__(self, **models: Model) -> None:
+        r"""Initialize a `ModelDict`.
+
+        Args:
+            models: An arbitrary number of models. Each model can be any type
+                of BoTorch `Model`, including multi-output models and `ModelList`.
+        """
+        if any(not isinstance(m, Model) for m in models.values()):
+            raise InputDataError(
+                f"Expected all models to be a BoTorch `Model`. Got {models}."
+            )
+        super().__init__(modules=models)

test/acquisition/test_acquisition.py

@@ -8,8 +8,10 @@
 from botorch.acquisition.acquisition import (
     AcquisitionFunction,
     MCSamplerMixin,
+    MultiModelAcquisitionFunction,
     OneShotAcquisitionFunction,
 )
+from botorch.models.model import ModelDict
 from botorch.sampling.normal import IIDNormalSampler
 from botorch.sampling.stochastic_samplers import StochasticSampler
 from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
@@ -25,6 +27,11 @@ def forward(self, X):
         raise NotImplementedError
 
 
+class DummyMultiModelAcqf(MultiModelAcquisitionFunction):
+    def forward(self, X):
+        raise NotImplementedError
+
+
 class TestAcquisitionFunction(BotorchTestCase):
     def test_abstract_raises(self):
         with self.assertRaises(TypeError):
@@ -48,3 +55,15 @@ def test_mc_sampler_mixin(self):
         sampler = IIDNormalSampler(sample_shape=torch.Size([2]))
         acqf.sampler = sampler
         self.assertIs(acqf.sampler, sampler)
+
+
+class TestMultiModelAcquisitionFunction(BotorchTestCase):
+    def test_multi_model_acquisition_function(self):
+        model_dict = ModelDict(
+            m1=MockModel(MockPosterior()),
+            m2=MockModel(MockPosterior()),
+        )
+        with self.assertRaises(TypeError):
+            MultiModelAcquisitionFunction(model_dict=model_dict)
+        acqf = DummyMultiModelAcqf(model_dict=model_dict)
+        self.assertIs(acqf.model_dict, model_dict)

test/models/test_model.py

@@ -8,12 +8,13 @@
 
 import torch
 from botorch.acquisition.objective import PosteriorTransform
+from botorch.exceptions.errors import InputDataError
 from botorch.models.deterministic import GenericDeterministicModel
-from botorch.models.model import Model, ModelList
+from botorch.models.model import Model, ModelDict, ModelList
 from botorch.models.utils import parse_training_data
 from botorch.posteriors.deterministic import DeterministicPosterior
 from botorch.posteriors.posterior_list import PosteriorList
-from botorch.utils.testing import BotorchTestCase
+from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
 
 
 class NotSoAbstractBaseModel(Model):
@@ -117,3 +118,15 @@ def test_posterior_transform(self):
                 posterior_tf.mean, torch.cat((2 * m1(X) + 1, 2 * m2(X) + 1), dim=-1)
             )
         )
+
+
+class TestModelDict(BotorchTestCase):
+    def test_model_dict(self):
+        models = {"m1": MockModel(MockPosterior()), "m2": MockModel(MockPosterior())}
+        model_dict = ModelDict(**models)
+        self.assertIs(model_dict["m1"], models["m1"])
+        self.assertIs(model_dict["m2"], models["m2"])
+        with self.assertRaisesRegex(
+            InputDataError, "Expected all models to be a BoTorch `Model`."
+        ):
+            ModelDict(m=MockPosterior())