SVDKL for Regression

[tkelestemur]

Hello all,

I'm trying to apply the SVDKL method (as shown in this example) to a regression problem. I have a pre-trained neural network that takes a latent vector (dim=256) and a (x,y,z) point. So the input to the neural network is a Nx259 and it is trained to regress SDF values of the 3d points. The network consists of basically a 8 fully connected layers.

Following the SVDKL example, I removed the laster linear layer of the network. This way it outputs a 512-dimensional feature vector.
Then I changed the Softmax likelihood with an GaussianLikelihood. The training code and the model definition can be seen here.
When I try to train the network, I get the following error:

Traceback (most recent call last):
  File "train_deep_gpis.py", line 210, in <module>
    main(conf)
  File "train_deep_gpis.py", line 95, in main
    train_loss = train_batch(
  File "train_deep_gpis.py", line 42, in train_batch
    loss = -loss_fn(output, sdf)
  File "/home/tarik/venvs/vtsc/lib/python3.8/site-packages/gpytorch/module.py", line 30, in __call__
    outputs = self.forward(*inputs, **kwargs)
  File "/home/tarik/venvs/vtsc/lib/python3.8/site-packages/gpytorch/mlls/variational_elbo.py", line 77, in forward
    return super().forward(variational_dist_f, target, **kwargs)
  File "/home/tarik/venvs/vtsc/lib/python3.8/site-packages/gpytorch/mlls/_approximate_mll.py", line 57, in forward
    log_likelihood = self._log_likelihood_term(approximate_dist_f, target, **kwargs).div(num_batch)
  File "/home/tarik/venvs/vtsc/lib/python3.8/site-packages/gpytorch/mlls/variational_elbo.py", line 61, in _log_likelihood_term
    return self.likelihood.expected_log_prob(target, variational_dist_f, **kwargs).sum(-1)
  File "/home/tarik/venvs/vtsc/lib/python3.8/site-packages/gpytorch/likelihoods/gaussian_likelihood.py", line 43, in expected_log_prob
    noise = noise.view(*noise.shape[:-1], *input.event_shape)
RuntimeError: shape '[2048, 2048, 512]' is invalid for input of size 1048576

Then I changed my likelihood to the MultitaskGaussianLikelihood. This seems to train but now my output is a vector-valued and I can't regress the underlying problem. I know the original DKL has an example for regression but my dataset is in the orders of a few million so I can't train it without mini-batch sampling.