Fix fit_gpytorch_model for batched models with custom modules (#193)

Summary:
Previously model fitting with sequential=True would break for custom batched models due to shortcomings of the converter.

This makes sure to catch the respective errors in the fitting and resort to fitting with sequential=False in this case.
Pull Request resolved: #193

Test Plan: unit tests

Reviewed By: sdaulton

Differential Revision: D16034229

Pulled By: Balandat

fbshipit-source-id: f9b42a7dbe94e0377dec17b66da7ebb2195337ef

Author: Balandat

botorch/fit.py

@@ -14,14 +14,20 @@
 from gpytorch.mlls.marginal_log_likelihood import MarginalLogLikelihood
 from gpytorch.mlls.sum_marginal_log_likelihood import SumMarginalLogLikelihood
 
-from .exceptions.warnings import OptimizationWarning
+from .exceptions.errors import UnsupportedError
+from .exceptions.warnings import BotorchWarning, OptimizationWarning
 from .models.converter import batched_to_model_list, model_list_to_batched
-from .models.gp_regression import HeteroskedasticSingleTaskGP
 from .models.gpytorch import BatchedMultiOutputGPyTorchModel
 from .optim.fit import fit_gpytorch_scipy
 from .optim.utils import sample_all_priors
 
 
+FAILED_CONVERSION_MSG = (
+    "Failed to convert ModelList to batched model. "
+    "Performing joint instead of sequential fitting."
+)
+
+
 def fit_gpytorch_model(
     mll: MarginalLogLikelihood, optimizer: Callable = fit_gpytorch_scipy, **kwargs: Any
 ) -> MarginalLogLikelihood:
@@ -60,20 +66,26 @@ def fit_gpytorch_model(
         isinstance(mll.model, BatchedMultiOutputGPyTorchModel)
         and mll.model._num_outputs > 1
         and sequential
-        and not isinstance(mll.model, HeteroskedasticSingleTaskGP)
     ):
-        model_list = batched_to_model_list(mll.model)
-        mll_ = SumMarginalLogLikelihood(model_list.likelihood, model_list)
-        fit_gpytorch_model(
-            mll=mll_,
-            optimizer=optimizer,
-            sequential=True,
-            max_retries=max_retries,
-            **kwargs,
-        )
-        model_ = model_list_to_batched(mll_.model)
-        mll.model.load_state_dict(model_.state_dict())
-        mll.eval()
+        try:  # check if backwards-conversion is possible
+            model_list = batched_to_model_list(mll.model)
+            model_ = model_list_to_batched(model_list)
+            mll_ = SumMarginalLogLikelihood(model_list.likelihood, model_list)
+            fit_gpytorch_model(
+                mll=mll_,
+                optimizer=optimizer,
+                sequential=True,
+                max_retries=max_retries,
+                **kwargs,
+            )
+            model_ = model_list_to_batched(mll_.model)
+            mll.model.load_state_dict(model_.state_dict())
+            return mll.eval()
+        except (NotImplementedError, UnsupportedError, RuntimeError, AttributeError):
+            warnings.warn(FAILED_CONVERSION_MSG, BotorchWarning)
+            return fit_gpytorch_model(
+                mll=mll, optimizer=optimizer, sequential=False, max_retries=max_retries
+            )
     # retry with random samples from the priors upon failure
     mll.train()
     original_state_dict = deepcopy(mll.model.state_dict())
@@ -91,5 +103,4 @@ def fit_gpytorch_model(
             logging.warning(f"Fitting failed on try {retry}.")
 
     warnings.warn("Fitting failed on all retries.", OptimizationWarning)
-    mll.eval()
-    return mll
+    return mll.eval()

botorch/models/converter.py

@@ -61,6 +61,12 @@ def _check_compatibility(models: ModelListGP) -> None:
             "Conversion of HeteroskedasticSingleTaskGP is currently unsupported."
         )
 
+    # TODO: Add support for custom likelihoods
+    if any(hasattr(m, "_likelihood_state_dict") for m in models):
+        raise NotImplementedError(
+            "Conversion of models with custom likelihoods is currently unsupported."
+        )
+
     # check that each model is single-output
     if not all(m._num_outputs == 1 for m in models):
         raise UnsupportedError("All models must be single-output.")
@@ -158,6 +164,11 @@ def batched_to_model_list(batch_model: BatchedMultiOutputGPyTorchModel) -> Model
         >>> batch_gp = SingleTaskGP(train_X, train_Y)
         >>> list_gp = batched_to_model_list(batch_gp)
     """
+    # TODO: Add support for HeteroskedasticSingleTaskGP
+    if isinstance(batch_model, HeteroskedasticSingleTaskGP):
+        raise NotImplementedError(
+            "Conversion of HeteroskedasticSingleTaskGP currently not supported."
+        )
     batch_sd = batch_model.state_dict()
 
     tensors = {n for n, p in batch_sd.items() if len(p.shape) > 0}
@@ -177,11 +188,6 @@ def batched_to_model_list(batch_model: BatchedMultiOutputGPyTorchModel) -> Model
         if isinstance(batch_model, FixedNoiseGP):
             noise_covar = batch_model.likelihood.noise_covar
             kwargs["train_Yvar"] = noise_covar.noise.select(input_bdims, i).clone()
-        # TODO: Add support for HeteroskedasticSingleTaskGP
-        if isinstance(batch_model, HeteroskedasticSingleTaskGP):
-            raise NotImplementedError(
-                "Conversion of HeteroskedasticSingleTaskGP currently not supported."
-            )
         model = batch_model.__class__(**kwargs)
         model.load_state_dict(sd)
         models.append(model)

test/models/test_converter.py

@@ -13,6 +13,7 @@
     SingleTaskGP,
 )
 from botorch.models.converter import batched_to_model_list, model_list_to_batched
+from gpytorch.likelihoods import GaussianLikelihood
 
 from .test_gpytorch import SimpleGPyTorchModel
 
@@ -93,6 +94,10 @@ def test_model_list_to_batched(self, cuda=False):
             gp2 = HeteroskedasticSingleTaskGP(
                 train_X, train_Y1, torch.ones_like(train_Y1)
             )
+            with self.assertRaises(NotImplementedError):
+                model_list_to_batched(ModelListGP(gp2))
+            # test custom likelihood
+            gp2 = SingleTaskGP(train_X, train_Y2, likelihood=GaussianLikelihood())
             with self.assertRaises(NotImplementedError):
                 model_list_to_batched(ModelListGP(gp2))
             # test FixedNoiseGP

test/models/test_gp_regression.py

@@ -98,7 +98,7 @@ def test_gp(self, cuda=False):
                     pvar = posterior_pred.variance
                     pvar_exp = _get_pvar_expected(posterior, model, X, num_outputs)
                     self.assertTrue(
-                        torch.allclose(pvar, pvar_exp, rtol=1e-4, atol=1e-06)
+                        torch.allclose(pvar, pvar_exp, rtol=1e-4, atol=1e-05)
                     )
 
                     # test batch evaluation
@@ -119,7 +119,7 @@ def test_gp(self, cuda=False):
                     pvar = posterior_pred.variance
                     pvar_exp = _get_pvar_expected(posterior, model, X, num_outputs)
                     self.assertTrue(
-                        torch.allclose(pvar, pvar_exp, rtol=1e-4, atol=1e-06)
+                        torch.allclose(pvar, pvar_exp, rtol=1e-4, atol=1e-05)
                     )
 
     def test_gp_cuda(self):