fit_gyptorch_model refactor (#1134)

Summary:
X-link: facebook/Ax#1134

Pull Request resolved: #1371

This commit updates `fit_gpytorch_model` and related methods, with the aim of fixing existing issues and improving extensibility. Key changes are as follow:

- Replace `fit_gpytorch_model` with `fit_gpytorch_mll`, a `Dispatcher` backed reimplementation of the original model fitting pipeline. Note that `fit_gpytorch_mll` does **not** pass `kwargs` to `optimizer` and instead introduces an optional `optimizer_kwargs` argument.

- Convert `fit_gpytorch_model` into a convenience method for calling `fit_gpytorch_mll` with (limited) support for legacy API.

- Add validation for multioutput GP fitting routines based on decomposing a single model into a list of independent models.

- Updated unit tests for relevant code paths.

Reviewed By: Balandat

Differential Revision: D38692173

fbshipit-source-id: 828cff264715cfa84ca4c4361db434574cf8fbf5

Author: James Wilson

README.md

@@ -126,7 +126,7 @@ For more details see our [Documentation](https://botorch.org/docs/introduction)
   ```python
   import torch
   from botorch.models import SingleTaskGP
-  from botorch.fit import fit_gpytorch_model
+  from botorch.fit import fit_gpytorch_mll
   from gpytorch.mlls import ExactMarginalLogLikelihood
 
   train_X = torch.rand(10, 2)
@@ -136,7 +136,7 @@ For more details see our [Documentation](https://botorch.org/docs/introduction)
 
   gp = SingleTaskGP(train_X, train_Y)
   mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
-  fit_gpytorch_model(mll)
+  fit_gpytorch_mll(mll)
   ```
 
 2. Construct an acquisition function

botorch/__init__.py

@@ -14,7 +14,11 @@
     test_functions,
 )
 from botorch.cross_validation import batch_cross_validation
-from botorch.fit import fit_fully_bayesian_model_nuts, fit_gpytorch_model
+from botorch.fit import (
+    fit_fully_bayesian_model_nuts,
+    fit_gpytorch_mll,
+    fit_gpytorch_model,
+)
 from botorch.generation.gen import (
     gen_candidates_scipy,
     gen_candidates_torch,
@@ -34,6 +38,7 @@
     "batch_cross_validation",
     "exceptions",
     "fit_fully_bayesian_model_nuts",
+    "fit_gpytorch_mll",
     "fit_gpytorch_model",
     "gen_candidates_scipy",
     "gen_candidates_torch",

botorch/cross_validation.py

@@ -13,7 +13,7 @@
 from typing import Any, Dict, NamedTuple, Optional, Type
 
 import torch
-from botorch.fit import fit_gpytorch_model
+from botorch.fit import fit_gpytorch_mll
 from botorch.models.gpytorch import GPyTorchModel
 from botorch.optim.utils import _filter_kwargs
 from botorch.posteriors.gpytorch import GPyTorchPosterior
@@ -119,7 +119,7 @@ def batch_cross_validation(
             internally. Note: Multi-task GPs are not currently supported.
         mll_cls: A MarginalLogLikelihood class.
         cv_folds: A CVFolds tuple.
-        fit_args: Arguments passed along to fit_gpytorch_model
+        fit_args: Arguments passed along to fit_gpytorch_mll.
 
     Returns:
         A CVResults tuple with the following fields
@@ -153,7 +153,7 @@ def batch_cross_validation(
     model_cv = model_cls(**_filter_kwargs(model_cls, **kwargs))
     mll_cv = mll_cls(model_cv.likelihood, model_cv)
     mll_cv.to(cv_folds.train_X)
-    mll_cv = fit_gpytorch_model(mll_cv, **fit_args)
+    mll_cv = fit_gpytorch_mll(mll_cv, **fit_args)
 
     # Evaluate on the hold-out set in batch mode
     with torch.no_grad():

botorch/exceptions/errors.py

@@ -37,3 +37,9 @@ class BotorchTensorDimensionError(BotorchError):
     r"""Exception raised when a tensor violates a botorch convention."""
 
     pass
+
+
+class ModelFittingError(Exception):
+    r"""Exception raised when attempts to fit a model terminate unsuccessfully."""
+
+    pass