Add num_outputs property to the Model API (#330)

Summary:
Pull Request resolved: #330

Addresses #295.

Reviewed By: liangshi7

Differential Revision: D18737863

fbshipit-source-id: 3df4d6b2b2b65609ac508e990d735838df33357b

Author: Balandat

botorch/models/cost.py

@@ -47,6 +47,7 @@ def __init__(
         self.fixed_cost = fixed_cost
         weights = torch.tensor([fidelity_weights[i] for i in self.fidelity_dims])
         self.register_buffer("weights", weights)
+        self._num_outputs = 1
 
     def forward(self, X: Tensor) -> Tensor:
         r"""Evaluate the cost on a candidate set X.

botorch/models/deterministic.py

@@ -37,6 +37,11 @@ def forward(self, X: Tensor) -> Tensor:
         """
         pass  # pragma: no cover
 
+    @property
+    def num_outputs(self) -> int:
+        r"""The number of outputs of the model."""
+        return self._num_outputs
+
     def posterior(
         self, X: Tensor, output_indices: Optional[List[int]] = None, **kwargs: Any
     ) -> DeterministicPosterior:
@@ -55,16 +60,18 @@ def posterior(
 class GenericDeterministicModel(DeterministicModel):
     r"""A generic deterministic model constructed from a callable."""
 
-    def __init__(self, f: Callable[[Tensor], Tensor]) -> None:
+    def __init__(self, f: Callable[[Tensor], Tensor], num_outputs: int = 1) -> None:
         r"""A generic deterministic model constructed from a callable.
 
         Args:
             f: A callable mapping a `batch_shape x n x d`-dim input tensor `X`
                 to a `batch_shape x n x m`-dimensional output tensor (the
                 outcome dimension `m` must be explicit, even if `m=1`).
+            num_outputs: The number of outputs `m`.
         """
         super().__init__()
         self._f = f
+        self._num_outputs = num_outputs
 
     def forward(self, X: Tensor) -> Tensor:
         r"""Compute the (deterministic) model output at X.
@@ -104,6 +111,7 @@ def __init__(self, a: Tensor, b: Union[Tensor, float] = 0.01) -> None:
         super().__init__()
         self.register_buffer("a", a)
         self.register_buffer("b", b.expand(a.size(-1)))
+        self._num_outputs = a.size(-1)
 
     def forward(self, X: Tensor) -> Tensor:
         return self.b + torch.einsum("...d,dm", X, self.a)

botorch/models/gpytorch.py

@@ -12,7 +12,7 @@
 """
 
 import warnings
-from abc import ABC, abstractproperty
+from abc import ABC
 from typing import Any, Iterator, List, Optional, Tuple, Union
 
 import torch
@@ -89,6 +89,11 @@ def _validate_tensor_args(
                 f" {Yvar.shape})."
             )
 
+    @property
+    def num_outputs(self) -> int:
+        r"""The number of outputs of the model."""
+        return self._num_outputs
+
     def posterior(
         self, X: Tensor, observation_noise: Union[bool, Tensor] = False, **kwargs: Any
     ) -> GPyTorchPosterior:
@@ -361,11 +366,6 @@ class ModelListGPyTorchModel(GPyTorchModel, ABC):
     evaluation of submodels.
     """
 
-    @abstractproperty
-    def num_outputs(self) -> int:
-        r"""The number of outputs of the model."""
-        pass  # pragma: no cover
-
     def posterior(
         self,
         X: Tensor,

botorch/models/model.py

@@ -49,6 +49,12 @@ def posterior(
         """
         pass  # pragma: no cover
 
+    @property
+    def num_outputs(self) -> int:
+        r"""The number of outputs of the model."""
+        cls_name = self.__class__.__name__
+        raise NotImplementedError(f"{cls_name} does not define num_outputs property")
+
     def condition_on_observations(self, X: Tensor, Y: Tensor, **kwargs: Any) -> "Model":
         r"""Condition the model on new observations.
 

botorch/models/multitask.py

@@ -86,6 +86,7 @@ def __init__(
             if any(t not in all_tasks for t in output_tasks):
                 raise RuntimeError("All output tasks must be present in input data.")
         self._output_tasks = output_tasks
+        self._num_outputs = len(output_tasks)
 
         # TODO (T41270962): Support task-specific noise levels in likelihood
         likelihood = GaussianLikelihood(noise_prior=GammaPrior(1.1, 0.05))

test/models/test_cost.py

@@ -18,6 +18,7 @@ def test_affine_fidelity_cost_model(self):
                 X = torch.rand(*batch_shape, 3, 4, device=self.device, dtype=dtype)
                 # test default parameters
                 model = AffineFidelityCostModel()
+                self.assertEqual(model.num_outputs, 1)
                 self.assertEqual(model.fidelity_dims, [-1])
                 self.assertEqual(model.fixed_cost, 0.01)
                 cost = model(X)

test/models/test_deterministic.py

@@ -25,6 +25,7 @@ def f(X):
             return X.mean(dim=-1, keepdim=True)
 
         model = GenericDeterministicModel(f)
+        self.assertEqual(model.num_outputs, 1)
         X = torch.rand(3, 2)
         # basic test
         p = model.posterior(X)
@@ -34,7 +35,8 @@ def f(X):
         with self.assertRaises(UnsupportedError):
             model.posterior(X, observation_noise=True)
         # check output indices
-        model = GenericDeterministicModel(lambda X: X)
+        model = GenericDeterministicModel(lambda X: X, num_outputs=2)
+        self.assertEqual(model.num_outputs, 2)
         p = model.posterior(X, output_indices=[0])
         self.assertTrue(torch.equal(p.mean, X[..., [0]]))
 
@@ -48,6 +50,7 @@ def test_AffineDeterministicModel(self):
         # test one-dim output
         a = torch.rand(3, 1)
         model = AffineDeterministicModel(a)
+        self.assertEqual(model.num_outputs, 1)
         for shape in ((4, 3), (1, 4, 3)):
             X = torch.rand(*shape)
             p = model.posterior(X)
@@ -56,6 +59,7 @@ def test_AffineDeterministicModel(self):
         # # test two-dim output
         a = torch.rand(3, 2)
         model = AffineDeterministicModel(a)
+        self.assertEqual(model.num_outputs, 2)
         for shape in ((4, 3), (1, 4, 3)):
             X = torch.rand(*shape)
             p = model.posterior(X)

test/models/test_gpytorch.py

@@ -40,6 +40,7 @@ def __init__(self, train_X, train_Y, outcome_transform=None):
         self.covar_module = ScaleKernel(RBFKernel())
         if outcome_transform is not None:
             self.outcome_transform = outcome_transform
+        self._num_outputs = 1
         self.to(train_X)
 
     def forward(self, x):
@@ -84,6 +85,7 @@ def test_gpytorch_model(self):
             train_Y = torch.sin(train_X)
             # basic test
             model = SimpleGPyTorchModel(train_X, train_Y, octf)
+            self.assertEqual(model.num_outputs, 1)
             test_X = torch.rand(2, 1, **tkwargs)
             posterior = model.posterior(test_X)
             self.assertIsInstance(posterior, GPyTorchPosterior)
@@ -175,6 +177,7 @@ def test_batched_multi_output_gpytorch_model(self):
             train_Y = torch.cat([torch.sin(train_X), torch.cos(train_X)], dim=-1)
             # basic test
             model = SimpleBatchedMultiOutputGPyTorchModel(train_X, train_Y)
+            self.assertEqual(model.num_outputs, 2)
             test_X = torch.rand(2, 1, **tkwargs)
             posterior = model.posterior(test_X)
             self.assertIsInstance(posterior, GPyTorchPosterior)
@@ -226,6 +229,7 @@ def test_model_list_gpytorch_model(self):
             m1 = SimpleGPyTorchModel(train_X1, train_Y1)
             m2 = SimpleGPyTorchModel(train_X2, train_Y2)
             model = SimpleModelListGPyTorchModel(m1, m2)
+            self.assertEqual(model.num_outputs, 2)
             test_X = torch.rand(2, 1, **tkwargs)
             posterior = model.posterior(test_X)
             self.assertIsInstance(posterior, GPyTorchPosterior)

test/models/test_model.py

@@ -22,3 +22,5 @@ def test_not_so_abstract_base_model(self):
         model = NotSoAbstractBaseModel()
         with self.assertRaises(NotImplementedError):
             model.condition_on_observations(None, None)
+        with self.assertRaises(NotImplementedError):
+            model.num_outputs

test/models/test_multitask.py

@@ -65,13 +65,11 @@ def _get_fixed_noise_model_single_output(**tkwargs):
 
 class TestMultiTaskGP(BotorchTestCase):
     def test_MultiTaskGP(self):
-        for double in (False, True):
-            tkwargs = {
-                "device": self.device,
-                "dtype": torch.double if double else torch.float,
-            }
+        for dtype in (torch.float, torch.double):
+            tkwargs = {"device": self.device, "dtype": dtype}
             model = _get_model(**tkwargs)
             self.assertIsInstance(model, MultiTaskGP)
+            self.assertEqual(model.num_outputs, 2)
             self.assertIsInstance(model.likelihood, GaussianLikelihood)
             self.assertIsInstance(model.mean_module, ConstantMean)
             self.assertIsInstance(model.covar_module, ScaleKernel)
@@ -140,13 +138,11 @@ def test_MultiTaskGP(self):
                 model.posterior(test_x)
 
     def test_MultiTaskGP_single_output(self):
-        for double in (False, True):
-            tkwargs = {
-                "device": self.device,
-                "dtype": torch.double if double else torch.float,
-            }
+        for dtype in (torch.float, torch.double):
+            tkwargs = {"device": self.device, "dtype": dtype}
             model = _get_model_single_output(**tkwargs)
             self.assertIsInstance(model, MultiTaskGP)
+            self.assertEqual(model.num_outputs, 1)
             self.assertIsInstance(model.likelihood, GaussianLikelihood)
             self.assertIsInstance(model.mean_module, ConstantMean)
             self.assertIsInstance(model.covar_module, ScaleKernel)
@@ -180,13 +176,11 @@ def test_MultiTaskGP_single_output(self):
 
 class TestFixedNoiseMultiTaskGP(BotorchTestCase):
     def test_FixedNoiseMultiTaskGP(self):
-        for double in (False, True):
-            tkwargs = {
-                "device": self.device,
-                "dtype": torch.double if double else torch.float,
-            }
+        for dtype in (torch.float, torch.double):
+            tkwargs = {"device": self.device, "dtype": dtype}
             model = _get_fixed_noise_model(**tkwargs)
             self.assertIsInstance(model, FixedNoiseMultiTaskGP)
+            self.assertEqual(model.num_outputs, 2)
             self.assertIsInstance(model.likelihood, FixedNoiseGaussianLikelihood)
             self.assertIsInstance(model.mean_module, ConstantMean)
             self.assertIsInstance(model.covar_module, ScaleKernel)
@@ -253,13 +247,11 @@ def test_FixedNoiseMultiTaskGP(self):
                 FixedNoiseMultiTaskGP(train_X, train_Y, train_Yvar, 0, output_tasks=[2])
 
     def test_FixedNoiseMultiTaskGP_single_output(self):
-        for double in (False, True):
-            tkwargs = {
-                "device": self.device,
-                "dtype": torch.double if double else torch.float,
-            }
+        for dtype in (torch.float, torch.double):
+            tkwargs = {"device": self.device, "dtype": dtype}
             model = _get_fixed_noise_model_single_output(**tkwargs)
             self.assertIsInstance(model, FixedNoiseMultiTaskGP)
+            self.assertEqual(model.num_outputs, 1)
             self.assertIsInstance(model.likelihood, FixedNoiseGaussianLikelihood)
             self.assertIsInstance(model.mean_module, ConstantMean)
             self.assertIsInstance(model.covar_module, ScaleKernel)