Repository: cornellius-gp/gpytorch

Issue:
[Bug] Exact DKL (Deep Kernel Learning) Regression w/ KISS-GP: only Tensors of floating point dtype can require gradients
# 🐛 Bug

I am trying to use Exact Deep Kernel Learning with KISS-GP, following the example from the [docs](https://docs.gpytorch.ai/en/stable/examples/06_PyTorch_NN_Integration_DKL/KISSGP_Deep_Kernel_Regression_CUDA.html#Overview). I am running into the error error `only Tensors of floating point dtype can require gradients` when I try to train using the Adam optimizer. It looks like it comes during back propagation through the Cholesky decomposition. I put in a print statement in [_pivoted_cholesky.py](https://github.com/cornellius-gp/gpytorch/blob/master/gpytorch/functions/_pivoted_cholesky.py#L94-L101) to print out `_matrix_args`
```python
...
    def backward(ctx, grad_output, _):
        full_permutation, short_permutation, *_matrix_args = ctx.saved_tensors
        _inverse_permutation = inverse_permutation(full_permutation)
        m = short_permutation.size(-1)  # The rank of the pivoted Cholesky factor

        with torch.enable_grad():
            # Create a new set of matrix args that we can backpropagate through
            [print(m, m.dtype, m.shape) for m in _matrix_args]
            matrix_args = [matrix_arg.detach().requires_grad_(True) for matrix_arg in _matrix_args]
```
which shows the output:
```
tensor([0.6931, 0.6725, 0.6141, 0.5279, 0.4271, 0.3252, 0.2331, 0.1573, 0.0999,
        0.0597, 0.0336, 0.0178, 0.0089, 0.0042, 0.0018, 0.0008],
       dtype=torch.float64, grad_fn=<ExpandBackward0>) torch.float64 torch.Size([16])
tensor([0.6931, 0.6761, 0.6274, 0.5539, 0.4653, 0.3718, 0.2827, 0.2045, 0.1407,
        0.0921, 0.0574, 0.0340, 0.0192, 0.0103, 0.0052, 0.0025],
       dtype=torch.float64, grad_fn=<ExpandBackward0>) torch.float64 torch.Size([16])
tensor([[ 85,  86,  87,  ..., 134, 135, 136],
        [103, 104, 105,  ..., 152, 153, 154],
        [134, 135, 136,  ..., 183, 184, 185],
        ...,
        [136, 137, 138,  ..., 185, 186, 187],
        [ 85,  86,  87,  ..., 134, 135, 136],
        [138, 139, 140,  ..., 187, 188, 189]]) torch.int64 torch.Size([30944, 16])
tensor([[ 1.1035e-03, -9.7422e-03, -1.5181e-02,  ..., -3.0474e-02,
         -4.7487e-02,  4.9617e-03],
        [ 3.3665e-03, -3.3772e-02, -1.7328e-02,  ..., -4.5008e-03,
         -2.3093e-03,  2.5992e-04],
        [ 9.1307e-04, -2.5656e-02, -1.4260e-03,  ..., -7.0527e-02,
         -3.9200e-03,  2.2810e-04],
        ...,
        [ 1.7785e-03, -7.3554e-02, -2.3944e-03,  ..., -3.6257e-02,
         -1.1803e-03,  4.9666e-05],
        [ 2.5766e-04, -1.4614e-02, -3.1947e-04,  ..., -6.2168e-02,
         -1.3590e-03,  4.3256e-05],
        [ 6.9518e-05, -9.5303e-04, -1.9168e-04,  ..., -1.9286e-02,
         -3.8789e-03,  3.7815e-04]], dtype=torch.float64,
       grad_fn=<ExpandBackward0>) torch.float64 torch.Size([30944, 16])
tensor([[ 85,  86,  87,  ..., 134, 135, 136],
        [103, 104, 105,  ..., 152, 153, 154],
        [134, 135, 136,  ..., 183, 184, 185],
        ...,
        [136, 137, 138,  ..., 185, 186, 187],
        [ 85,  86,  87,  ..., 134, 135, 136],
        [138, 139, 140,  ..., 187, 188, 189]]) torch.int64 torch.Size([30944, 16])
tensor([[ 1.1035e-03, -9.7422e-03, -1.5181e-02,  ..., -3.0474e-02,
         -4.7487e-02,  4.9617e-03],
        [ 3.3665e-03, -3.3772e-02, -1.7328e-02,  ..., -4.5008e-03,
         -2.3093e-03,  2.5992e-04],
        [ 9.1307e-04, -2.5656e-02, -1.4260e-03,  ..., -7.0527e-02,
         -3.9200e-03,  2.2810e-04],
        ...,
        [ 1.7785e-03, -7.3554e-02, -2.3944e-03,  ..., -3.6257e-02,
         -1.1803e-03,  4.9666e-05],
        [ 2.5766e-04, -1.4614e-02, -3.1947e-04,  ..., -6.2168e-02,
         -1.3590e-03,  4.3256e-05],
        [ 6.9518e-05, -9.5303e-04, -1.9168e-04,  ..., -1.9286e-02,
         -3.8789e-03,  3.7815e-04]], dtype=torch.float64,
       grad_fn=<ExpandBackward0>) torch.float64 torch.Size([30944, 16])
```

You can see there are two tensors of size [30944, 16] that have integer dtype being used. 16 is the grid size for the Grid Interpolation Kernel, so I think it has something to do with how that kernel is computing the covariance matrix. I don't know what those matrices are though.

## To reproduce

** Code snippet to reproduce **

Here is my code minus data loading.

```python
import gpytorch
import torch
class LargeFeatureExtractor(torch.nn.Sequential):
    def __init__(self, dim, latent_dim):
        super(LargeFeatureExtractor, self).__init__()
        self.latent_dim = latent_dim
        self.dim = dim
        self.add_module('linear1', torch.nn.Linear(dim, 1000))
        self.add_module('relu1', torch.nn.ReLU())
        self.add_module('linear2', torch.nn.Linear(1000, 500))
        self.add_module('relu2', torch.nn.ReLU())
        self.add_module('linear3', torch.nn.Linear(500, 50))
        self.add_module('relu3', torch.nn.ReLU())
        self.add_module('linear4', torch.nn.Linear(50, latent_dim))

class DeepKernelGP(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood, feature_extractor, grid_size):
        super(DeepKernelGP, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.GridInterpolationKernel(
            gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(ard_num_dims=feature_extractor.latent_dim)),
            num_dims=feature_extractor.latent_dim, grid_size=grid_size
        )
        
        self.feature_extractor = feature_extractor

        # This module will scale the NN features so that they're nice values
        self.scale_to_bounds = gpytorch.utils.grid.ScaleToBounds(-1., 1.)

    def forward(self, x):
        # We're first putting our data through a deep net (feature extractor)
        projected_x = self.feature_extractor(x)
        projected_x = self.scale_to_bounds(projected_x)  # Make the NN values "nice"

        mean_x = self.mean_module(projected_x)
        covar_x = self.covar_module(projected_x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

# I have a data set X of shape [30944, 14] and y of shape [30944, 1]
likelihood = gpytorch.likelihoods.GaussianLikelihood(
     noise_prior=gpytorch.priors.SmoothedBoxPrior(0.1, 1.5, sigma=0.001)
)
latent_dim = 2
grid_size = 16
feature_extractor = LargeFeatureExtractor(X.shape[1], latent_dim)
model = DeepKernelGP(X, y, likelihood, feature_extractor, grid_size)
loss = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)
optimizer = torch.optim.Adam(
    [
        {'params': model.feature_extractor.parameters()},
        {'params': model.covar_module.parameters()},
        {'params': model.mean_module.parameters()},
        {'params': model.likelihood.parameters()},
    ], 
    lr=args.lr
)

def closure():
    optimizer.zero_grad()
    output = model(x)
    nll = -loss(output, y)
    nll.backward()
    return nll

while True:
    optimizer.step(closure=closure)
```

** Stack trace/error message **
```
Traceback (most recent call last):
  File "/vast/home/stetzler/nuclei_summer/nuclei/nuclei/train/train.py", line 87, in _train
    curr_loss = optimizer.step(closure=closure)
  File "/vast/home/stetzler/.conda/envs/nuclei/lib/python3.10/site-packages/torch/optim/optimizer.py", line 109, in wrapper
    return func(*args, **kwargs)
  File "/vast/home/stetzler/.conda/envs/nuclei/lib/python3.10/site-packages/torch/autograd/grad_mode.py", line 27, in decorate_context
    return func(*args, **kwargs)
  File "/vast/home/stetzler/.conda/envs/nuclei/lib/python3.10/site-packages/torch/optim/adam.py", line 118, in step
    loss = closure()
  File "/vast/home/stetzler/nuclei_summer/nuclei/nuclei/train/train.py", line 58, in closure
    nll.backward()
  File "/vast/home/stetzler/.conda/envs/nuclei/lib/python3.10/site-packages/torch/_tensor.py", line 396, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "/vast/home/stetzler/.conda/envs/nuclei/lib/python3.10/site-packages/torch/autograd/__init__.py", line 173, in backward
    Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/vast/home/stetzler/.conda/envs/nuclei/lib/python3.10/site-packages/torch/autograd/function.py", line 253, in apply
    return user_fn(self, *args)
  File "/vast/home/stetzler/.conda/envs/nuclei/lib/python3.10/site-packages/gpytorch/functions/_pivoted_cholesky.py", line 102, in backward
    matrix_args = [matrix_arg.detach().requires_grad_(True) for matrix_arg in _matrix_args]
  File "/vast/home/stetzler/.conda/envs/nuclei/lib/python3.10/site-packages/gpytorch/functions/_pivoted_cholesky.py", line 102, in <listcomp>
    matrix_args = [matrix_arg.detach().requires_grad_(True) for matrix_arg in _matrix_args]
RuntimeError: only Tensors of floating point dtype can require gradients```
```

## Expected Behavior

I expect the training to run.

## System information

**Please complete the following information:**
- <!-- GPyTorch Version (run `print(gpytorch.__version__)` --> python -c "import gpytorch; print(gpytorch.__version__)"
1.7.0
- <!-- PyTorch Version (run `print(torch.__version__)` --> python -c "import torch; print(torch.__version__)"
1.12.0+cu102
- <!-- Computer OS --> cat /etc/os-release | head -n 2
NAME="Red Hat Enterprise Linux"
VERSION="8.4 (Ootpa)"

## Additional context
Add any other context about the problem here.



Discussion:
I have encountered the same issue. A workaround for now is to delete the GridInterpolationKernel, resulting in same similar MAE.
Thanks, yes my workaround is to use the `InducingPointKernel` instead of the `GridInterpolationKernel` paired with the neural network to maintain speed up with many data points.
> Thanks, yes my workaround is to use the `InducingPointKernel` instead of the `GridInterpolationKernel` paired with the neural network to maintain speed up with many data points.

Greetings,

I also met the same problem. I tried to change `GridInterpolationKernel` to `InducingPointKernel`, but I can't figure out the dimensions. It throws an error and says `"x1_ and x2_ must have the same dimension"`. Could you please let me know how can I successfully change the kernel to `InducingPointKernel`?
> Could you please let me know how can I successfully change the kernel to `InducingPointKernel`?

Here is code that works for my application
```
class LargeFeatureExtractor(torch.nn.Sequential):
    def __init__(self, dim, latent_dim):
        super(LargeFeatureExtractor, self).__init__()
        self.latent_dim = latent_dim
        self.dim = dim
        self.add_module('linear1', torch.nn.Linear(dim, 1000))
        self.add_module('relu1', torch.nn.ReLU())
        self.add_module('linear2', torch.nn.Linear(1000, 500))
        self.add_module('relu2', torch.nn.ReLU())
        self.add_module('linear3', torch.nn.Linear(500, 50))
        self.add_module('relu3', torch.nn.ReLU())
        self.add_module('linear4', torch.nn.Linear(50, latent_dim))

class DeepKernelSparseGP(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood, inducing_points):
        super(DeepKernelSparseGP, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.base_covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(ard_num_dims=feature_extractor.latent_dim))       
        self.covar_module = gpytorch.kernels.InducingPointKernel(
            self.base_covar_module, inducing_points=inducing_points, likelihood=likelihood
        )

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

class _DeepKernelSparseGP():
    def __init__(self, gp, feature_extractor):
        self.gp = gp
        self.feature_extractor = feature_extractor
        self.scale_to_bounds = gpytorch.utils.grid.ScaleToBounds(-1., 1.)
    
    def __call__(self, x):
        return self.gp(self.scale_to_bounds(self.feature_extractor(x)))

    def train(self):
        return self.gp.train()

    def eval(self):
        return self.gp.eval()

# data are [N, d] where N=number of samples, d=number of dimensions
# data are scaled to [-1, 1]
# args.latent_dim is number of dimensions of the neutral network / GP latent space
# args.num_points is the number of inducing points
feature_extractor = LargeFeatureExtractor(data['train']['x'].shape[1], args.latent_dim)
inducing_points = torch.rand(size=(args.num_points, args.latent_dim)) * 2 - 1
likelihood = gpytorch.likelihoods.GaussianLikelihood()
gp = DeepKernelSparseGP(data['train']['x'], data['train']['y'], likelihood, inducing_points)
model = _DeepKernelSparseGP(gp, feature_extractor)
```
> > Could you please let me know how can I successfully change the kernel to `InducingPointKernel`?
> 
> Here is code that works for my application
> 
> ```
> class LargeFeatureExtractor(torch.nn.Sequential):
>     def __init__(self, dim, latent_dim):
>         super(LargeFeatureExtractor, self).__init__()
>         self.latent_dim = latent_dim
>         self.dim = dim
>         self.add_module('linear1', torch.nn.Linear(dim, 1000))
>         self.add_module('relu1', torch.nn.ReLU())
>         self.add_module('linear2', torch.nn.Linear(1000, 500))
>         self.add_module('relu2', torch.nn.ReLU())
>         self.add_module('linear3', torch.nn.Linear(500, 50))
>         self.add_module('relu3', torch.nn.ReLU())
>         self.add_module('linear4', torch.nn.Linear(50, latent_dim))
> 
> class DeepKernelSparseGP(gpytorch.models.ExactGP):
>     def __init__(self, train_x, train_y, likelihood, inducing_points):
>         super(DeepKernelSparseGP, self).__init__(train_x, train_y, likelihood)
>         self.mean_module = gpytorch.means.ConstantMean()
>         self.base_covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(ard_num_dims=feature_extractor.latent_dim))       
>         self.covar_module = gpytorch.kernels.InducingPointKernel(
>             self.base_covar_module, inducing_points=inducing_points, likelihood=likelihood
>         )
> 
>     def forward(self, x):
>         mean_x = self.mean_module(x)
>         covar_x = self.covar_module(x)
>         return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
> 
> class _DeepKernelSparseGP():
>     def __init__(self, gp, feature_extractor):
>         self.gp = gp
>         self.feature_extractor = feature_extractor
>         self.scale_to_bounds = gpytorch.utils.grid.ScaleToBounds(-1., 1.)
>     
>     def __call__(self, x):
>         return self.gp(self.scale_to_bounds(self.feature_extractor(x)))
> 
>     def train(self):
>         return self.gp.train()
> 
>     def eval(self):
>         return self.gp.eval()
> 
> # data are [N, d] where N=number of samples, d=number of dimensions
> # data are scaled to [-1, 1]
> # args.latent_dim is number of dimensions of the neutral network / GP latent space
> # args.num_points is the number of inducing points
> feature_extractor = LargeFeatureExtractor(data['train']['x'].shape[1], args.latent_dim)
> inducing_points = torch.rand(size=(args.num_points, args.latent_dim)) * 2 - 1
> likelihood = gpytorch.likelihoods.GaussianLikelihood()
> gp = DeepKernelSparseGP(data['train']['x'], data['train']['y'], likelihood, inducing_points)
> model = _DeepKernelSparseGP(gp, feature_extractor)
> ```

Thank you so much! This helps me a lot.
Glad it is helpful. I found there is a mistake in that code, change
```
        self.base_covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(ard_num_dims=feature_extractor.latent_dim))       
```
to
```
        self.base_covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(ard_num_dims=inducing_points.shape[1]))       
```

this provides an RBF length scale in each dimension of the latent space
Hi, sorry that code I shared actually does not work at all. This will work for `InducingPointKernel`:

```
class LargeFeatureExtractor(torch.nn.Sequential):
    def __init__(self, dim, latent_dim):
        super(LargeFeatureExtractor, self).__init__()
        self.latent_dim = latent_dim
        self.dim = dim
        self.add_module('linear1', torch.nn.Linear(dim, 1000))
        self.add_module('relu1', torch.nn.ReLU())
        self.add_module('linear2', torch.nn.Linear(1000, 500))
        self.add_module('relu2', torch.nn.ReLU())
        self.add_module('linear3', torch.nn.Linear(500, 50))
        self.add_module('relu3', torch.nn.ReLU())
        self.add_module('linear4', torch.nn.Linear(50, latent_dim))

class DeepKernelSparseGP(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood, feature_extractor, inducing_points):
        super(DeepKernelSparseGP, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.base_covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(ard_num_dims=inducing_points.shape[1]))       
        self.covar_module = gpytorch.kernels.InducingPointKernel(
            self.base_covar_module, inducing_points=inducing_points, likelihood=likelihood
        )
        self.scale_to_bounds = gpytorch.utils.grid.ScaleToBounds(-1., 1.)
        self.feature_extractor = feature_extractor

    def forward(self, x):
        x = self.scale_to_bounds(self.feature_extractor(x))
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

```

What I was doing before, splitting the GP and FeatureExtractor classes was necessary when using `gpytorch.models.ApproximateGP` for using stochastic variational inference. For SVI, the following classes work for me:

```
class LargeFeatureExtractor(torch.nn.Sequential):
    def __init__(self, dim, latent_dim):
        super(LargeFeatureExtractor, self).__init__()
        self.latent_dim = latent_dim
        self.dim = dim
        self.add_module('linear1', torch.nn.Linear(dim, 1000))
        self.add_module('relu1', torch.nn.ReLU())
        self.add_module('linear2', torch.nn.Linear(1000, 500))
        self.add_module('relu2', torch.nn.ReLU())
        self.add_module('linear3', torch.nn.Linear(500, 50))
        self.add_module('relu3', torch.nn.ReLU())
        self.add_module('linear4', torch.nn.Linear(50, latent_dim))

class _DeepKernelSVGP():
    def __init__(self, gp, feature_extractor):
        self.gp = gp
        self.feature_extractor = feature_extractor
        self.scale_to_bounds = gpytorch.utils.grid.ScaleToBounds(-1., 1.)
    
    def __call__(self, x):
        return self.gp(self.scale_to_bounds(self.feature_extractor(x)))

    def train(self):
        return self.gp.train()

    def eval(self):
        return self.gp.eval()

class DeepKernelSVGP(gpytorch.models.ApproximateGP):
    def __init__(self, likelihood, feature_extractor, inducing_points):
        variational_distribution = gpytorch.variational.TrilNaturalVariationalDistribution(inducing_points.size(0))
        variational_strategy = gpytorch.variational.VariationalStrategy(
            self, inducing_points, variational_distribution, learn_inducing_locations=True
        )
        super(DeepKernelSVGP, self).__init__(variational_strategy)
        self.likelihood = likelihood
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.ScaleKernel(gpytorch.kernels.RBFKernel(ard_num_dims=feature_extractor.latent_dim))

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

```
You really saved my day. I did get errors with the previous code, I fixed some but others kept coming up. This one works for me. I see that the problem is I should put the call to ` feature_extractor` inside the forward function of the GP Model class. Thank you so much! I really appreciate your help
Thanks @stevenstetzler and @rottenivy for pointing this out - I think that this bug escaped our unit tests for some reason. I'll put up a PR to fix the error with `GridInterpolationKernel`, though switching to `InducingPointKernel` is also a fine workaround for right now.

Resolving patch:
diff --git a/gpytorch/functions/_pivoted_cholesky.py b/gpytorch/functions/_pivoted_cholesky.py
--- a/gpytorch/functions/_pivoted_cholesky.py
+++ b/gpytorch/functions/_pivoted_cholesky.py
@@ -98,7 +98,11 @@ def backward(ctx, grad_output, _):
 
         with torch.enable_grad():
             # Create a new set of matrix args that we can backpropagate through
-            matrix_args = [matrix_arg.detach().requires_grad_(True) for matrix_arg in _matrix_args]
+            matrix_args = []
+            for matrix_arg in _matrix_args:
+                if matrix_arg.dtype in (torch.float, torch.double, torch.half):
+                    matrix_arg = matrix_arg.detach().requires_grad_(True)
+                matrix_args.append(matrix_arg)
 
             # Create new linear operator using new matrix args
             matrix = ctx.representation_tree(*matrix_args)
diff --git a/gpytorch/lazy/lazy_tensor.py b/gpytorch/lazy/lazy_tensor.py
--- a/gpytorch/lazy/lazy_tensor.py
+++ b/gpytorch/lazy/lazy_tensor.py
@@ -1651,7 +1651,8 @@ def _set_requires_grad(self, val):
                     arg.requires_grad_(val)
         for arg in self._kwargs.values():
             if hasattr(arg, "requires_grad"):
-                arg.requires_grad_(val)
+                if arg.dtype in (torch.float, torch.double, torch.half):
+                    arg.requires_grad_(val)
 
     @requires_grad.setter
     def requires_grad(self, val):


Test patch: