Fix flaky SVGP classification test (#2495)

Author: gpleiss

gpytorch/test/base_keops_test_case.py

@@ -66,7 +66,7 @@ def test_forward_x1_neq_x2(self, use_keops=True, ard=False, **kwargs):
                 # The patch makes sure that we're actually using KeOps
                 k1 = kern1(x1, x2).to_dense()
                 k2 = kern2(x1, x2).to_dense()
-                self.assertLess(torch.norm(k1 - k2), 1e-4)
+                self.assertLess(torch.norm(k1 - k2), 1e-3)
 
         if use_keops:
             self.assertTrue(keops_mock.called)
@@ -86,7 +86,7 @@ def test_batch_matmul(self, use_keops=True, **kwargs):
                 # The patch makes sure that we're actually using KeOps
                 res1 = kern1(x1, x1).matmul(rhs)
                 res2 = kern2(x1, x1).matmul(rhs)
-                self.assertLess(torch.norm(res1 - res2), 1e-4)
+                self.assertLess(torch.norm(res1 - res2), 1e-3)
 
         if use_keops:
             self.assertTrue(keops_mock.called)
@@ -115,7 +115,7 @@ def test_gradient(self, use_keops=True, ard=False, **kwargs):
                 # stack all gradients into a tensor
                 grad_s1 = torch.vstack(torch.autograd.grad(s1, [*kern1.hyperparameters()]))
                 grad_s2 = torch.vstack(torch.autograd.grad(s2, [*kern2.hyperparameters()]))
-                self.assertAllClose(grad_s1, grad_s2, rtol=1e-4, atol=1e-5)
+                self.assertAllClose(grad_s1, grad_s2, rtol=1e-3, atol=1e-3)
 
         if use_keops:
             self.assertTrue(keops_mock.called)

setup.py

@@ -39,6 +39,7 @@ def find_version(*file_paths):
 
 torch_min = "1.11"
 install_requires = [
+    "mpmath>=0.19,<=1.3",  # avoid incompatibiltiy with torch+sympy with mpmath 1.4
     "scikit-learn",
     "scipy",
     "linear_operator>=0.5.2",

test/examples/test_svgp_gp_classification.py

@@ -16,7 +16,7 @@
 
 
 def train_data(cuda=False):
-    train_x = torch.linspace(0, 1, 260)
+    train_x = torch.linspace(0, 1, 150)
     train_y = torch.cos(train_x * (2 * math.pi)).gt(0).float()
     if cuda:
         return train_x.cuda(), train_y.cuda()
@@ -49,7 +49,7 @@ class TestSVGPClassification(BaseTestCase, unittest.TestCase):
     def test_classification_error(self, cuda=False, mll_cls=gpytorch.mlls.VariationalELBO):
         train_x, train_y = train_data(cuda=cuda)
         likelihood = BernoulliLikelihood()
-        model = SVGPClassificationModel(torch.linspace(0, 1, 25))
+        model = SVGPClassificationModel(torch.linspace(0, 1, 64))
         mll = mll_cls(likelihood, model, num_data=len(train_y))
         if cuda:
             likelihood = likelihood.cuda()
@@ -59,12 +59,12 @@ def test_classification_error(self, cuda=False, mll_cls=gpytorch.mlls.Variationa
         # Find optimal model hyperparameters
         model.train()
         likelihood.train()
-        optimizer = optim.Adam([{"params": model.parameters()}, {"params": likelihood.parameters()}], lr=0.1)
+        optimizer = optim.Adam([{"params": model.parameters()}, {"params": likelihood.parameters()}], lr=0.03)
 
         _wrapped_cg = MagicMock(wraps=linear_operator.utils.linear_cg)
         _cg_mock = patch("linear_operator.utils.linear_cg", new=_wrapped_cg)
         with _cg_mock as cg_mock:
-            for _ in range(400):
+            for _ in range(100):
                 optimizer.zero_grad()
                 output = model(train_x)
                 loss = -mll(output, train_y)