Repository: cornellius-gp/gpytorch
Issue: [Bug] Reloading saved parameters into a variational model hurts performance

# 🐛 Bug

The performance of the saved SVGP model differs a lot if the model is not newly instantiated. 

## To reproduce

** Code snippet to reproduce **
```python
import os

from math import floor

import tqdm
import gpytorch
from gpytorch.models import ApproximateGP
from gpytorch.variational import CholeskyVariationalDistribution
from gpytorch.variational import VariationalStrategy
from scipy.io import loadmat
from sklearn.metrics import mean_squared_error
import torch
from torch.utils.data import TensorDataset, DataLoader
import urllib.request


if not os.path.isfile('../elevators.mat'):
    print('Downloading \\'elevators\\' UCI dataset...')
    urllib.request.urlretrieve('https://drive.google.com/uc?export=download&id=1jhWL3YUHvXIaftia4qeAyDwVxo6j1alk', '../elevators.mat')

data = torch.Tensor(loadmat('../elevators.mat')['data'])
X = data[:, :-1]
X = X - X.min(0)[0]
X = 2 * (X / X.max(0)[0]) - 1
y = data[:, -1]


train_n = int(floor(0.8 * len(X)))
train_x = X[:train_n, :].contiguous()
train_y = y[:train_n].contiguous()

test_x = X[train_n:, :].contiguous()
test_y = y[train_n:].contiguous()


train_dataset = TensorDataset(train_x, train_y)
train_loader = DataLoader(train_dataset, batch_size=1024, shuffle=True)

test_dataset = TensorDataset(test_x, test_y)
test_loader = DataLoader(test_dataset, batch_size=1024, shuffle=False)


class GPModel(ApproximateGP):
    def __init__(self, inducing_points):
        variational_distribution = CholeskyVariationalDistribution(inducing_points.size(0))
        variational_strategy = VariationalStrategy(self, inducing_points, variational_distribution, learn_inducing_locations=True)
        super(GPModel, self).__init__(variational_strategy)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel())

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)


inducing_points = train_x[:500, :]
model = GPModel(inducing_points=inducing_points)
likelihood = gpytorch.likelihoods.GaussianLikelihood()

num_epochs = 4

model.train()
likelihood.train()

# We use SGD here, rather than Adam. Emperically, we find that SGD is better for variational regression
optimizer = torch.optim.Adam([
    {'params': model.parameters()},
    {'params': likelihood.parameters()},
], lr=0.01)

# Our loss object. We're using the VariationalELBO
mll = gpytorch.mlls.VariationalELBO(likelihood, model, num_data=train_y.size(0))


epochs_iter = tqdm.tqdm(range(num_epochs), desc="Epoch")
for i in epochs_iter:
    # Within each iteration, we will go over each minibatch of data
    minibatch_iter = tqdm.tqdm(train_loader, desc="Minibatch", leave=False)
    for x_batch, y_batch in minibatch_iter:
        optimizer.zero_grad()
        output = model(x_batch)
        loss = -mll(output, y_batch)
        minibatch_iter.set_postfix(loss=loss.item())
        loss.backward()
        optimizer.step()


model.eval()
likelihood.eval()
y_pred_list = []
y_true_list = []
with torch.no_grad():
    for x_batch, y_batch in test_loader:
        preds = model(x_batch)
        y_pred_list.append(preds.mean)
        y_true_list.append(y_batch)

y_pred = torch.cat(y_pred_list, dim=0).numpy()
y_true = torch.cat(y_true_list, dim=0).numpy()

test_mse = mean_squared_error(y_true, y_pred)
print(test_mse)


# torch.save(model.state_dict(), 'my_gp_with_nn_model.pth')    # uncomment once to save the model on the disk
state_dict = torch.load('my_gp_with_nn_model.pth')

# inducing_points = train_x[:500, :]
# model = GPModel(inducing_points=inducing_points)
# likelihood = gpytorch.likelihoods.GaussianLikelihood()

model.load_state_dict(state_dict)

model.eval()
likelihood.eval()
y_pred_list = []
y_true_list = []
with torch.no_grad():
    for x_batch, y_batch in test_loader:
        preds = model(x_batch)
        y_pred_list.append(preds.mean)
        y_true_list.append(y_batch)

y_pred = torch.cat(y_pred_list, dim=0).numpy()
y_true = torch.cat(y_true_list, dim=0).numpy()

test_mse = mean_squared_error(y_true, y_pred)
print(test_mse)
```

** Stack trace/error message **
```
0.010551954
0.23515734   # or 0.5873753 or 0.40703666, unpredictable
```

## Expected Behavior

The performance of a saved model is expected to be consistent after loading from the saved file. However, whether instantiate it from class seems to affect the performance a lot. 

In the snippet, the interesting part is 

```python
inducing_points = train_x[:500, :]
model = GPModel(inducing_points=inducing_points)
likelihood = gpytorch.likelihoods.GaussianLikelihood()
```
If I uncomment it, i.e., I instantiate a new model and load the weights to it, the performance is always consistent. The test MSE of the trained model and saved model is almost very close to each other. 
However, if I don't instantiate a new model, but rather load the weights to an existing trained model, the performance is unpredictable. 

## System information

**Please complete the following information:**
- GPyTorch Version: 1.2.0
- PyTorch Version: 1.6.0 
- Debian 10

## Additional context
I came across the bug when I saved the model to disk based on some validation score and reload it after the whole training and test it on a test set in the same script. The reloading was done to an existing trained model, rather than a newly instantiated model. When I start a new script, instantiate a new model, load the saved weights to double-check the test performance, the number turns out to be different. 

I also tested the same thing on ExactGP, which doesn't have such a problem. I am not sure whether such behavior is expected by the library itself. Pure PyTorch modules don't seem to have such behavior. 

Changed files:
gpytorch/variational/_variational_strategy.py, gpytorch/kernels/inducing_point_kernel.py, gpytorch/module.py, gpytorch/kernels/grid_kernel.py, gpytorch/models/exact_gp.py

diff --git a/gpytorch/kernels/grid_kernel.py b/gpytorch/kernels/grid_kernel.py
index a514247e..a275ba51 100644
--- a/gpytorch/kernels/grid_kernel.py
+++ b/gpytorch/kernels/grid_kernel.py
@@ -61,16 +61,15 @@ class GridKernel(Kernel):
         if not self.interpolation_mode:
             self.register_buffer("full_grid", create_data_from_grid(grid))
 
+    def _clear_cache(self):
+        if hasattr(self, "_cached_kernel_mat"):
+            del self._cached_kernel_mat
+
     def register_buffer_list(self, base_name, tensors):
         """Helper to register several buffers at once under a single base name"""
         for i, tensor in enumerate(tensors):
             self.register_buffer(base_name + "_" + str(i), tensor)
 
-    def train(self, mode=True):
-        if hasattr(self, "_cached_kernel_mat"):
-            del self._cached_kernel_mat
-        return super(GridKernel, self).train(mode)
-
     @property
     def grid(self):
         return [getattr(self, f"grid_{i}") for i in range(self.num_dims)]
@@ -91,8 +90,7 @@ class GridKernel(Kernel):
         if not self.interpolation_mode:
             self.full_grid = create_data_from_grid(self.grid)
 
-        if hasattr(self, "_cached_kernel_mat"):
-            del self._cached_kernel_mat
+        self._clear_cache()
         return self
 
     @property
diff --git a/gpytorch/kernels/inducing_point_kernel.py b/gpytorch/kernels/inducing_point_kernel.py
index 4f1b73ff..fa6fedc6 100644
--- a/gpytorch/kernels/inducing_point_kernel.py
+++ b/gpytorch/kernels/inducing_point_kernel.py
@@ -25,10 +25,9 @@ class InducingPointKernel(Kernel):
         self.register_parameter(name="inducing_points", parameter=torch.nn.Parameter(inducing_points))
         self.register_added_loss_term("inducing_point_loss_term")
 
-    def train(self, mode=True):
+    def _clear_cache(self):
         if hasattr(self, "_cached_kernel_mat"):
             del self._cached_kernel_mat
-        return super(InducingPointKernel, self).train(mode)
 
     @property
     def _inducing_mat(self):
diff --git a/gpytorch/models/exact_gp.py b/gpytorch/models/exact_gp.py
index 036fdebe..2c0c0228 100644
--- a/gpytorch/models/exact_gp.py
+++ b/gpytorch/models/exact_gp.py
@@ -85,6 +85,10 @@ class ExactGP(GP):
             self.train_targets = fn(self.train_targets)
         return super(ExactGP, self)._apply(fn)
 
+    def _clear_cache(self):
+        # The precomputed caches from test time live in prediction_strategy
+        self.prediction_strategy = None
+
     def local_load_samples(self, samples_dict, memo, prefix):
         """
         Replace the model's learned hyperparameters with samples from a posterior distribution.
@@ -236,19 +240,6 @@ class ExactGP(GP):
 
         return new_model
 
-    def train(self, mode=True):
-        if mode:
-            self.prediction_strategy = None
-        return super(ExactGP, self).train(mode)
-
-    def _load_from_state_dict(
-        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
-    ):
-        self.prediction_strategy = None
-        super()._load_from_state_dict(
-            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
-        )
-
     def __call__(self, *args, **kwargs):
         train_inputs = list(self.train_inputs) if self.train_inputs is not None else []
         inputs = [i.unsqueeze(-1) if i.ndimension() == 1 else i for i in args]
diff --git a/gpytorch/module.py b/gpytorch/module.py
index 96ba783e..214c57b4 100644
--- a/gpytorch/module.py
+++ b/gpytorch/module.py
@@ -30,6 +30,13 @@ class Module(nn.Module):
             return [_validate_module_outputs(output) for output in outputs]
         return _validate_module_outputs(outputs)
 
+    def _clear_cache(self):
+        """
+        Clear any precomputed caches.
+        Should be implemented by any module that caches any computation at test time.
+        """
+        pass
+
     def _get_module_and_name(self, parameter_name):
         """Get module and name from full parameter name."""
         module, name = parameter_name.split(".", 1)
@@ -260,6 +267,12 @@ class Module(nn.Module):
         if new_constraint.initial_value is not None:
             self.initialize(**{param_name: new_constraint.initial_value})
 
+    def train(self, mode=True):
+        # If we're going in training mode, we need to clear any pre-comptued caches from eval mode
+        if (self.training and not mode) or mode:
+            self._clear_cache()
+        return super().train(mode=mode)
+
     def constraint_for_parameter_name(self, param_name):
         base_module = self
         base_name = param_name
@@ -290,6 +303,15 @@ class Module(nn.Module):
                 if key in state_dict:
                     param.data = state_dict[key].data
 
+    def _load_from_state_dict(
+        self, state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+    ):
+        # If we're loading from a state dict, we need to clear any precomputed caches
+        self._clear_cache()
+        super()._load_from_state_dict(
+            state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
+        )
+
     def load_strict_shapes(self, value):
         def apply_fn(module):
             module._load_strict_shapes = value
diff --git a/gpytorch/variational/_variational_strategy.py b/gpytorch/variational/_variational_strategy.py
index bbba3229..28ff7ef3 100644
--- a/gpytorch/variational/_variational_strategy.py
+++ b/gpytorch/variational/_variational_strategy.py
@@ -35,6 +35,9 @@ class _VariationalStrategy(Module, ABC):
         self._variational_distribution = variational_distribution
         self.register_buffer("variational_params_initialized", torch.tensor(0))
 
+    def _clear_cache(self):
+        clear_cache_hook(self)
+
     def _expand_inputs(self, x, inducing_points):
         """
         Pre-processing step in __call__ to make x the same batch_shape as the inducing points
@@ -94,12 +97,6 @@ class _VariationalStrategy(Module, ABC):
             kl_divergence = torch.distributions.kl.kl_divergence(self.variational_distribution, self.prior_distribution)
         return kl_divergence
 
-    def train(self, mode=True):
-        # Make sure we are clearing the cache if we change modes
-        if (self.training and not mode) or mode:
-            clear_cache_hook(self)
-        return super().train(mode=mode)
-
     def __call__(self, x, prior=False):
         # If we're in prior mode, then we're done!
         if prior:
@@ -107,7 +104,7 @@ class _VariationalStrategy(Module, ABC):
 
         # Delete previously cached items from the training distribution
         if self.training:
-            clear_cache_hook(self)
+            self._clear_cache()
         # (Maybe) initialize variational distribution
         if not self.variational_params_initialized.item():
             prior_dist = self.prior_distribution