Repository: cornellius-gp/gpytorch

Issue:
[Bug] Performance varies according to the batch size
# 🐛 Bug

<!-- A clear and concise description of what the bug is. -->

The accuracy of the classifier trained with the codes (https://github.com/cornellius-gp/gpytorch/tree/master/examples/06_PyTorch_NN_Integration_DKL) varies according to the batch size. It should be consistent regardless of the batch size. Am I missing something?

## To reproduce

** Code snippet to reproduce **
```python
from torch.optim import SGD, Adam
from torch.optim.lr_scheduler import MultiStepLR
import torch.nn.functional as F
from torch import nn
import torch
import os
import torchvision.datasets as dset
import torchvision.transforms as transforms
import gpytorch
import math
import tqdm


normalize = transforms.Normalize(mean=[0.5071, 0.4867, 0.4408], std=[0.2675, 0.2565, 0.2761])
aug_trans = [transforms.RandomCrop(32, padding=4), transforms.RandomHorizontalFlip()]
common_trans = [transforms.ToTensor(), normalize]
train_compose = transforms.Compose(aug_trans + common_trans)
test_compose = transforms.Compose(common_trans)


dataset = "cifar10"


if ('CI' in os.environ):  # this is for running the notebook in our testing framework
    train_set = torch.utils.data.TensorDataset(torch.randn(8, 3, 32, 32), torch.rand(8).round().long())
    test_set = torch.utils.data.TensorDataset(torch.randn(4, 3, 32, 32), torch.rand(4).round().long())
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=4, shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=2, shuffle=False)
    num_classes = 2
elif dataset == 'cifar10':
    train_set = dset.CIFAR10('data', train=True, transform=train_compose, download=True)
    test_set = dset.CIFAR10('data', train=False, transform=test_compose)
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=256, shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=16, shuffle=False)
    num_classes = 10
elif dataset == 'cifar100':
    train_set = dset.CIFAR100('data', train=True, transform=train_compose, download=True)
    test_set = dset.CIFAR100('data', train=False, transform=test_compose)
    train_loader = torch.utils.data.DataLoader(train_set, batch_size=256, shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_set, batch_size=256, shuffle=False)
    num_classes = 100
else:
    raise RuntimeError('dataset must be one of "cifar100" or "cifar10"')

from densenet import DenseNet

class DenseNetFeatureExtractor(DenseNet):
    def forward(self, x):
        features = self.features(x)
        out = F.relu(features, inplace=True)
        out = F.avg_pool2d(out, kernel_size=self.avgpool_size).view(features.size(0), -1)
        return out

feature_extractor = DenseNetFeatureExtractor(block_config=(6, 6, 6), num_classes=num_classes)
num_features = feature_extractor.classifier.in_features

class GaussianProcessLayer(gpytorch.models.ApproximateGP):
    def __init__(self, num_dim, grid_bounds=(-10., 10.), grid_size=64):
        variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
            num_inducing_points=grid_size, batch_shape=torch.Size([num_dim])
        )

        # Our base variational strategy is a GridInterpolationVariationalStrategy,
        # which places variational inducing points on a Grid
        # We wrap it with a IndependentMultitaskVariationalStrategy so that our output is a vector-valued GP
        variational_strategy = gpytorch.variational.IndependentMultitaskVariationalStrategy(
            gpytorch.variational.GridInterpolationVariationalStrategy(
                self, grid_size=grid_size, grid_bounds=[grid_bounds],
                variational_distribution=variational_distribution,
            ), num_tasks=num_dim,
        )
        super().__init__(variational_strategy)

        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(
                lengthscale_prior=gpytorch.priors.SmoothedBoxPrior(
                    math.exp(-1), math.exp(1), sigma=0.1, transform=torch.exp
                )
            )
        )
        self.mean_module = gpytorch.means.ConstantMean()
        self.grid_bounds = grid_bounds

    def forward(self, x):
        mean = self.mean_module(x)
        covar = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean, covar)


class DKLModel(gpytorch.Module):
    def __init__(self, feature_extractor, num_dim, grid_bounds=(-10., 10.)):
        super(DKLModel, self).__init__()
        self.feature_extractor = feature_extractor
        self.gp_layer = GaussianProcessLayer(num_dim=num_dim, grid_bounds=grid_bounds)
        self.grid_bounds = grid_bounds
        self.num_dim = num_dim

    def forward(self, x):
        features = self.feature_extractor(x)
        features = gpytorch.utils.grid.scale_to_bounds(features, self.grid_bounds[0], self.grid_bounds[1])
        # This next line makes it so that we learn a GP for each feature
        features = features.transpose(-1, -2).unsqueeze(-1)
        res = self.gp_layer(features)
        return res

model = DKLModel(feature_extractor, num_dim=num_features)
likelihood = gpytorch.likelihoods.SoftmaxLikelihood(num_features=model.num_dim, num_classes=num_classes)


# If you run this example without CUDA, I hope you like waiting!
if torch.cuda.is_available():
    model = model.cuda()
    likelihood = likelihood.cuda()

# load trained model
state_dict = torch.load('dkl_cifar_checkpoint.dat')
model_state_dict = state_dict['model']
likelihood_state_dict = state_dict['likelihood']

model.load_state_dict(model_state_dict)
likelihood.load_state_dict(likelihood_state_dict)

def test():
    model.eval()
    likelihood.eval()

    correct = 0
    with torch.no_grad(), gpytorch.settings.num_likelihood_samples(16):
        for data, target in test_loader:
            if torch.cuda.is_available():
                data, target = data.cuda(), target.cuda()
            output = likelihood(model(data))  # This gives us 16 samples from the predictive distribution
            pred = output.probs.mean(0).argmax(-1)  # Taking the mean over all of the sample we've drawn
            correct += pred.eq(target.view_as(pred)).cpu().sum()
    print('Test set: Accuracy: {}/{} ({}%)'.format(
        correct, len(test_loader.dataset), 100. * correct / float(len(test_loader.dataset))
    ))

with gpytorch.settings.use_toeplitz(False):
    test()

```

## Expected Behavior

<!-- A clear and concise description of what you expected to happen. -->
When the batch size is set to 256, the accuracy on the test set is 90.9% as expected. However, if the batch size is set to 16, the accuracy drops to 74.8%.


## System information

**Please complete the following information:**
- GPyTorch Version: 1.3.0
- PyTorch Version: 1.8.0.dev20201117
- Ubuntu 18.04

## Additional context
Add any other context about the problem here.



Discussion:
> The accuracy of the classifier trained with the codes (https://github.com/cornellius-gp/gpytorch/tree/master/examples/06_PyTorch_NN_Integration_DKL) varies according to the batch size. It should be consistent regardless of the batch size. 

Are you referring to the training set or test set batch size? Changing the training set batch size will affect the optimization of the the neural network parameters. A smaller batch size makes the optimization more stochastic, which makes it harder to learn good NN and variational parameters.
@gpleiss I am referring to the test set batch size. When we change the batch size in `test_loader = torch.utils.data.DataLoader(test_set, batch_size=16, shuffle=False)`, then the prediction performance varies. 

Resolving patch:
diff --git a/gpytorch/utils/grid.py b/gpytorch/utils/grid.py
--- a/gpytorch/utils/grid.py
+++ b/gpytorch/utils/grid.py
@@ -1,15 +1,62 @@
 #!/usr/bin/env python3
 
 import math
+import warnings
 from typing import List, Tuple
 
 import torch
 
 
-def scale_to_bounds(x, lower_bound, upper_bound):
+class ScaleToBounds(torch.nn.Module):
     """
     Scale the input data so that it lies in between the lower and upper bounds.
 
+    In training (`self.train()`), this module adjusts the scaling factor to the minibatch of data.
+    During evaluation (`self.eval()`), this module uses the scaling factor from the previous minibatch of data.
+
+    :param float lower_bound: lower bound of scaled data
+    :param float upper_bound: upper bound of scaled data
+
+    Example:
+        >>> train_x = torch.randn(10, 5)
+        >>> module = gpytorch.utils.grid.ScaleToBounds(lower_bound=-1., upper_bound=1.)
+        >>>
+        >>> module.train()
+        >>> scaled_train_x = module(train_x)  # Data should be between -0.95 and 0.95
+        >>>
+        >>> module.eval()
+        >>> test_x = torch.randn(10, 5)
+        >>> scaled_test_x = module(test_x)  # Scaling is based on train_x
+    """
+
+    def __init__(self, lower_bound, upper_bound):
+        super().__init__()
+        self.lower_bound = float(lower_bound)
+        self.upper_bound = float(upper_bound)
+        self.register_buffer("min_val", torch.tensor(lower_bound))
+        self.register_buffer("max_val", torch.tensor(upper_bound))
+
+    def forward(self, x):
+        if self.training:
+            min_val = x.min()
+            max_val = x.max()
+            self.min_val.data = min_val
+            self.max_val.data = max_val
+        else:
+            min_val = self.min_val
+            max_val = self.max_val
+            # Clamp extreme values
+            x = x.clamp(min_val, max_val)
+
+        diff = max_val - min_val
+        x = (x - min_val) * (0.95 * (self.upper_bound - self.lower_bound) / diff) + 0.95 * self.lower_bound
+        return x
+
+
+def scale_to_bounds(x, lower_bound, upper_bound):
+    """
+    DEPRECATRED: Use :obj:`~gpytorch.utils.grid.ScaleToBounds` instead.
+
     :param x: the input data
     :type x: torch.Tensor (... x n x d)
     :param float lower_bound: lower bound of scaled data
@@ -17,6 +64,10 @@ def scale_to_bounds(x, lower_bound, upper_bound):
     :return: scaled data
     :rtype: torch.Tensor (... x n x d)
     """
+    warnings.warn(
+        "The `scale_to_bounds` method is deprecated. Use the `gpytorch.utils.grid.ScaleToBounds` module instead.",
+        DeprecationWarning,
+    )
     # Scale features so they fit inside grid bounds
     min_val = x.min()
     max_val = x.max()


Test patch: