Gpytorch: SpectralMixtureKernel initialize_from_data CUDA bug

[tzoiker]

🐛 Bug

When running initialize_from_data with CUDA tensors, in this line zero tensor is initialized without input device consideration, so that code fails later.

To reproduce

import torch, gpytorch
train_x = torch.tensor([1.0,2.0]).cuda()
train_y = torch.ones(2).cuda()
kernel = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4).cuda()
kernel.initialize_from_data(train_x, train_y)

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-1-419d92435b05> in <module>
      3 train_y = torch.ones(2).cuda()
      4 kernel = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4).cuda()
----> 5 kernel.initialize_from_data(train_x, train_y)

~/storage/miniconda2/envs/water/lib/python3.6/site-packages/gpytorch/kernels/spectral_mixture_kernel.py in initialize_from_data(self, train_x, train_y, **kwargs)
    166         )
    167         # Draw means from Unif(0, 0.5 / minimum distance between two points)
--> 168         self.raw_mixture_means.data.uniform_().mul_(0.5).div_(min_dist)
    169         self.raw_mixture_means.data = self.raw_mixture_means_constraint.inverse_transform(self.raw_mixture_means.data)
    170         # Mixture weights should be roughly the stdv of the y values divided by the number of mixtures

RuntimeError: expected backend CUDA and dtype Float but got backend CPU and dtype Float

Expected Behavior

I expected it to get device info from training data.

System information

Please complete the following information:

• GPyTorch 0.3.2
• PyTorch 1.1.0
• MacOS 10.14

[jacobrgardner]

You need to cast your kernel to the GPU with a .cuda() call too.

[tzoiker]

You need to cast your kernel to the GPU with a .cuda() call too.

@jacobrgardner, I had mistakes in my comment, but the original source of error didn't get anywhere anyway. Check updated snippet and, please, reopen.

[adam-rysanek]

I can confirm I'm getting the same error. I can't seem to get the Spectral Mixture Model processing via CUDA in the same way many of the other gpytorch examples are laid out. I get exactly the same error at initialization of the kernel, as shown by @tzoiker above. Attempting to cast the kernel to .cuda() does not get rid of the error. Perhaps our syntax is wrong?

To reproduce

n = 40
train_x = torch.zeros(pow(n, 2), 2)
for i in range(n):
    for j in range(n):
        train_x[i * n + j][0] = float(i) / (n-1)
        train_x[i * n + j][1] = float(j) / (n-1)
# True function is sin( 2*pi*(x0+x1))
train_x = train_x.float().cuda()
train_y = torch.sin((train_x[:, 0] + train_x[:, 1]) * (2 * math.pi)) + torch.randn_like(train_x[:, 0]).mul(0.01)
train_y = train_y.float().cuda()

class SpectralMixtureGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(SpectralMixtureGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4,ard_num_dims=2).cuda()
        self.covar_module.initialize_from_data(train_x, train_y)

    def forward(self,x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)


likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = SpectralMixtureGPModel(train_x, train_y, likelihood).cuda()

[adam-rysanek]

I have come up with a solution that seems to be working, but it seems to require a different convention of where to cast to .cuda() andcpu() compared to other examples. For example, the implementation below is a variation on the gptorch KISS-GP multidimensional GP regression examples.

Here's what I've implemented for the above solution

Training data initialization:
Note: Training data is not cast to .cuda()

n = 40
train_x = torch.zeros(pow(n, 2), 2)
for i in range(n):
    for j in range(n):
        train_x[i * n + j][0] = float(i) / (n-1)
        train_x[i * n + j][1] = float(j) / (n-1)
# True function is sin( 2*pi*(x0+x1))
train_x = train_x.float()
train_y = torch.sin((train_x[:, 0] + train_x[:, 1]) * (2 * math.pi)) + torch.randn_like(train_x[:, 0]).mul(0.01)
train_y = train_y.float()

Establish model
Note: Only the model is cast to .cuda()

class SpectralMixtureGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(SpectralMixtureGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()
        self.covar_module = gpytorch.kernels.SpectralMixtureKernel(num_mixtures=4,ard_num_dims=2)
        self.covar_module.initialize_from_data(train_x, train_y)

    def forward(self,x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)


likelihood = gpytorch.likelihoods.GaussianLikelihood()
model = SpectralMixtureGPModel(train_x, train_y, likelihood).cuda()

Train model
Note: Now that the model is casted to .cuda(), but the training data is not, we cast train_x and train_y to .cuda() where they are called in-situ:

# Find optimal model hyperparameters
model.train()
likelihood.train()

# Use the adam optimizer
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)

# "Loss" for GPs - the marginal log likelihood
mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

training_iter = 100
for i in range(training_iter):
    optimizer.zero_grad()
    output = model(train_x.cuda())
    loss = -mll(output, train_y.cuda())
    loss.backward()
    print('Iter %d/%d - Loss: %.3f' % (i + 1, training_iter, loss.item()))
    optimizer.step()

Test model
Note: text_x is cast to .cuda() in order to generate samples, pred_labels, but these are then cast back to .cpu() so that the rest of the analysis is handled without issue.

# Set model and likelihood into evaluation mode
model.eval()
likelihood.eval()

# Generate nxn grid of test points spaced on a grid of size 1/(n-1) in [0,1]x[0,1]
n = 10
test_x = torch.zeros(int(pow(n, 2)), 2)
for i in range(n):
    for j in range(n):
        test_x[i * n + j][0] = float(i) / (n-1)
        test_x[i * n + j][1] = float(j) / (n-1)

with torch.no_grad(), gpytorch.settings.fast_pred_var():
    observed_pred = likelihood(model(test_x.cuda()))
    pred_labels = (observed_pred.mean.view(n, n)).cpu()

# Calc abosolute error
test_y_actual = torch.sin(((test_x[:, 0] + test_x[:, 1]) * (2 * math.pi))).view(n, n)
delta_y = torch.abs(pred_labels - test_y_actual).detach().numpy()

# Define a plotting function
def ax_plot(f, ax, y_labels, title):
    im = ax.imshow(y_labels)
    ax.set_title(title)
    f.colorbar(im)

# Plot our predictive means
f, observed_ax = plt.subplots(1, 1, figsize=(4, 3))
ax_plot(f, observed_ax, pred_labels, 'Predicted Values (Likelihood)')

# Plot the true values
f, observed_ax2 = plt.subplots(1, 1, figsize=(4, 3))
ax_plot(f, observed_ax2, test_y_actual, 'Actual Values (Likelihood)')

# Plot the absolute errors
f, observed_ax3 = plt.subplots(1, 1, figsize=(4, 3))
ax_plot(f, observed_ax3, delta_y, 'Absolute Error Surface')

This solution works for me (and is lightning fast!), but is it right?

[jacobrgardner]

@tzoiker @adam-rysanek Sorry this took an embarrassingly long time to fix for how simple it was.

85a000b should resolve the issue internally.

[gpleiss]

looks like this should've been closed with 85a000b