Suppress warnings during testing (#191)

Summary:
Suppresses a number of OptimizationWarnings emitted during testing caused by having very small limits on the number of optimization iterations (for unit testing performance reasons).
Pull Request resolved: #191

Reviewed By: danielrjiang

Differential Revision: D16025104

Pulled By: Balandat

fbshipit-source-id: f1cf19d5d4c0d691126cbe1d60e888f4913f34a2

Author: Balandat

botorch/models/converter.py

@@ -61,10 +61,6 @@ def _check_compatibility(models: ModelListGP) -> None:
             "Conversion of HeteroskedasticSingleTaskGP is currently unsupported."
         )
 
-    # if the list has only one model, we can just return a copy of that
-    if len(models) == 1:
-        return deepcopy(models[0])
-
     # 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.")
@@ -96,6 +92,10 @@ def model_list_to_batched(model_list: ModelListGP) -> BatchedMultiOutputGPyTorch
     models = model_list.models
     _check_compatibility(models)
 
+    # if the list has only one model, we can just return a copy of that
+    if len(models) == 1:
+        return deepcopy(models[0])
+
     # construct inputs
     train_X = deepcopy(models[0].train_inputs[0])
     train_Y = torch.stack([m.train_targets.clone() for m in models], dim=-1)

test/models/test_gp_regression.py

@@ -4,9 +4,11 @@
 
 import math
 import unittest
+import warnings
 
 import torch
 from botorch import fit_gpytorch_model
+from botorch.exceptions.warnings import OptimizationWarning
 from botorch.models.gp_regression import (
     FixedNoiseGP,
     HeteroskedasticSingleTaskGP,
@@ -63,7 +65,9 @@ def test_gp(self, cuda=False):
                     mll = ExactMarginalLogLikelihood(model.likelihood, model).to(
                         **tkwargs
                     )
-                    fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
+                    with warnings.catch_warnings():
+                        warnings.filterwarnings("ignore", category=OptimizationWarning)
+                        fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
 
                     # test init
                     self.assertIsInstance(model.mean_module, ConstantMean)

test/models/test_model_list_gp_regression.py

@@ -4,9 +4,11 @@
 
 import math
 import unittest
+import warnings
 
 import torch
-from botorch import fit_gpytorch_model
+from botorch.exceptions.warnings import OptimizationWarning
+from botorch.fit import fit_gpytorch_model
 from botorch.models import ModelListGP
 from botorch.models.gp_regression import FixedNoiseGP, SingleTaskGP
 from botorch.posteriors import GPyTorchPosterior
@@ -70,11 +72,15 @@ def test_ModelListGP(self, cuda=False):
                 self.assertIsInstance(mll_, ExactMarginalLogLikelihood)
 
             # test model fitting (sequential)
-            mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
-            # test model fitting (joint)
-            mll = fit_gpytorch_model(
-                mll, options={"maxiter": 1}, max_retries=1, sequential=False
-            )
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=OptimizationWarning)
+                mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=OptimizationWarning)
+                # test model fitting (joint)
+                mll = fit_gpytorch_model(
+                    mll, options={"maxiter": 1}, max_retries=1, sequential=False
+                )
 
             # test posterior
             test_x = torch.tensor([[0.25], [0.75]], **tkwargs)
@@ -140,7 +146,9 @@ def test_ModelListGP_fixed_noise(self, cuda=False):
             mll = SumMarginalLogLikelihood(model.likelihood, model)
             for mll_ in mll.mlls:
                 self.assertIsInstance(mll_, ExactMarginalLogLikelihood)
-            mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=OptimizationWarning)
+                mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
 
             # test posterior
             test_x = torch.tensor([[0.25], [0.75]], **tkwargs)

test/models/test_multitask.py

@@ -4,9 +4,11 @@
 
 import math
 import unittest
+import warnings
 
 import torch
-from botorch import fit_gpytorch_model
+from botorch.exceptions.warnings import OptimizationWarning
+from botorch.fit import fit_gpytorch_model
 from botorch.models.multitask import FixedNoiseMultiTaskGP, MultiTaskGP
 from botorch.posteriors import GPyTorchPosterior
 from gpytorch.distributions import MultitaskMultivariateNormal, MultivariateNormal
@@ -82,7 +84,9 @@ def test_MultiTaskGP(self, cuda=False):
 
             # test model fitting
             mll = ExactMarginalLogLikelihood(model.likelihood, model)
-            mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=OptimizationWarning)
+                mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
 
             # test posterior
             test_x = torch.rand(2, 1, **tkwargs)
@@ -155,7 +159,9 @@ def test_MultiTaskGP_single_output(self, cuda=False):
 
             # test model fitting
             mll = ExactMarginalLogLikelihood(model.likelihood, model)
-            mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=OptimizationWarning)
+                mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
 
             # test posterior
             test_x = torch.rand(2, 1, **tkwargs)
@@ -197,7 +203,9 @@ def test_FixedNoiseMultiTaskGP(self, cuda=False):
 
             # test model fitting
             mll = ExactMarginalLogLikelihood(model.likelihood, model)
-            mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=OptimizationWarning)
+                mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
 
             # test posterior
             test_x = torch.rand(2, 1, **tkwargs)
@@ -268,7 +276,9 @@ def test_FixedNoiseMultiTaskGP_single_output(self, cuda=False):
 
             # test model fitting
             mll = ExactMarginalLogLikelihood(model.likelihood, model)
-            mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", category=OptimizationWarning)
+                mll = fit_gpytorch_model(mll, options={"maxiter": 1}, max_retries=1)
 
             # test posterior
             test_x = torch.rand(2, 1, **tkwargs)

test/optim/test_utils.py

@@ -3,6 +3,7 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
 import unittest
+import warnings
 from copy import deepcopy
 
 import torch
@@ -229,7 +230,10 @@ def test_sample_all_priors(self, cuda=False):
                 outputscale_prior=GammaPrior(2.0, 0.15),
             )
             original_state_dict = dict(deepcopy(mll.model.state_dict()))
-            sample_all_priors(model)
+            with warnings.catch_warnings(record=True) as ws:
+                sample_all_priors(model)
+                self.assertEqual(len(ws), 1)
+                self.assertTrue("rsample" in str(ws[0].message))
 
             # the lengthscale should not have changed because sampling is
             # not implemented for SmoothedBoxPrior

test/test_cross_validation.py

@@ -4,9 +4,11 @@
 
 import math
 import unittest
+import warnings
 
 import torch
 from botorch.cross_validation import batch_cross_validation, gen_loo_cv_folds
+from botorch.exceptions.warnings import OptimizationWarning
 from botorch.models.gp_regression import FixedNoiseGP, SingleTaskGP
 from gpytorch.mlls.exact_marginal_log_likelihood import ExactMarginalLogLikelihood
 
@@ -42,12 +44,14 @@ def test_single_task_batch_cv(self, cuda=False):
                         train_X=train_X, train_Y=train_Y
                     )
                     # Test SingleTaskGP
-                    cv_results = batch_cross_validation(
-                        model_cls=SingleTaskGP,
-                        mll_cls=ExactMarginalLogLikelihood,
-                        cv_folds=noiseless_cv_folds,
-                        fit_args={"options": {"maxiter": 1}},
-                    )
+                    with warnings.catch_warnings():
+                        warnings.filterwarnings("ignore", category=OptimizationWarning)
+                        cv_results = batch_cross_validation(
+                            model_cls=SingleTaskGP,
+                            mll_cls=ExactMarginalLogLikelihood,
+                            cv_folds=noiseless_cv_folds,
+                            fit_args={"options": {"maxiter": 1}},
+                        )
                     expected_shape = batch_shape + torch.Size([n, 1, num_outputs])
                     self.assertEqual(cv_results.posterior.mean.shape, expected_shape)
                     self.assertEqual(cv_results.observed_Y.shape, expected_shape)
@@ -56,12 +60,14 @@ def test_single_task_batch_cv(self, cuda=False):
                     noisy_cv_folds = gen_loo_cv_folds(
                         train_X=train_X, train_Y=train_Y, train_Yvar=train_Yvar
                     )
-                    cv_results = batch_cross_validation(
-                        model_cls=FixedNoiseGP,
-                        mll_cls=ExactMarginalLogLikelihood,
-                        cv_folds=noisy_cv_folds,
-                        fit_args={"options": {"maxiter": 1}},
-                    )
+                    with warnings.catch_warnings():
+                        warnings.filterwarnings("ignore", category=OptimizationWarning)
+                        cv_results = batch_cross_validation(
+                            model_cls=FixedNoiseGP,
+                            mll_cls=ExactMarginalLogLikelihood,
+                            cv_folds=noisy_cv_folds,
+                            fit_args={"options": {"maxiter": 1}},
+                        )
                     self.assertEqual(cv_results.posterior.mean.shape, expected_shape)
                     self.assertEqual(cv_results.observed_Y.shape, expected_shape)
                     self.assertEqual(cv_results.observed_Y.shape, expected_shape)

test/test_end_to_end.py

@@ -4,9 +4,11 @@
 
 import math
 import unittest
+import warnings
 
 import torch
 from botorch.acquisition import ExpectedImprovement, qExpectedImprovement
+from botorch.exceptions.warnings import OptimizationWarning
 from botorch.fit import fit_gpytorch_model
 from botorch.models import FixedNoiseGP, SingleTaskGP
 from botorch.optim import joint_optimize
@@ -35,19 +37,23 @@ def _setUp(self, double=False, cuda=False):
         self.mll_st = ExactMarginalLogLikelihood(
             self.model_st.likelihood, self.model_st
         )
-        self.mll_st = fit_gpytorch_model(
-            self.mll_st, options={"maxiter": 5}, max_retries=1
-        )
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=OptimizationWarning)
+            self.mll_st = fit_gpytorch_model(
+                self.mll_st, options={"maxiter": 5}, max_retries=1
+            )
         model_fn = FixedNoiseGP(
             self.train_x, self.train_y, self.train_yvar.expand_as(self.train_y)
         )
         self.model_fn = model_fn.to(device=device, dtype=dtype)
         self.mll_fn = ExactMarginalLogLikelihood(
             self.model_fn.likelihood, self.model_fn
         )
-        self.mll_fn = fit_gpytorch_model(
-            self.mll_fn, options={"maxiter": 5}, max_retries=1
-        )
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=OptimizationWarning)
+            self.mll_fn = fit_gpytorch_model(
+                self.mll_fn, options={"maxiter": 5}, max_retries=1
+            )
 
     def test_qEI(self, cuda=False):
         for double in (True, False):

test/test_fit.py

@@ -150,11 +150,8 @@ def test_fit_gpytorch_model_singular(self, cuda=False):
             gp = SingleTaskGP(X_train, Y_train, likelihood=test_likelihood)
             mll = ExactMarginalLogLikelihood(gp.likelihood, gp)
             mll.to(device=device, dtype=dtype)
-            with warnings.catch_warnings(record=True) as ws:
-                # this will do multiple retries
-                fit_gpytorch_model(mll, options=options, max_retries=1)
-                self.assertEqual(len(ws), 1)
-                self.assertTrue(MAX_RETRY_MSG in str(ws[0].message))
+            # this will do multiple retries (and emit warnings, which is desired)
+            fit_gpytorch_model(mll, options=options, max_retries=2)
 
     def test_fit_gpytorch_model_singular_cuda(self):
         if torch.cuda.is_available():

test/test_gen.py

@@ -4,9 +4,11 @@
 
 import math
 import unittest
+import warnings
 
 import torch
 from botorch.acquisition import qExpectedImprovement
+from botorch.exceptions.warnings import OptimizationWarning
 from botorch.fit import fit_gpytorch_model
 from botorch.gen import gen_candidates_scipy, gen_candidates_torch
 from botorch.models import SingleTaskGP
@@ -37,7 +39,11 @@ def _setUp(self, double=False, cuda=False, expand=False):
         model = SingleTaskGP(self.train_x, self.train_y)
         self.model = model.to(device=device, dtype=dtype)
         self.mll = ExactMarginalLogLikelihood(self.model.likelihood, self.model)
-        self.mll = fit_gpytorch_model(self.mll, options={"maxiter": 1}, max_retries=1)
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", category=OptimizationWarning)
+            self.mll = fit_gpytorch_model(
+                self.mll, options={"maxiter": 1}, max_retries=1
+            )
 
 
 class TestGenCandidates(TestBaseCandidateGeneration):