Why does CPU training work well, but GPU training doesn't?

[MJimitater]

Hi GPytorch Team, you've done great work, thanks for this library, I love it!
I have a question regarding using GPytorch with CPU, which works well in my case, and GPU, which somehow doesn't seem to work in my case.

Currently Im trying to get a Classifier to work.
Im working with a non-Gaussian likelihood, as in the tutorial.

I noticed how it works well on the CPU for a small subset of my large dataset:

Training on CPU

X_train.shape, y_train.shape, X_test.shape, y_test.shape
(torch.Size([320, 240]), torch.Size([320]), torch.Size([80, 240]), torch.Size([80]))

from gpytorch.models import ApproximateGP
from gpytorch.variational import CholeskyVariationalDistribution
from gpytorch.variational import VariationalStrategy
class GPClassificationModel(ApproximateGP):
    def __init__(self, train_x):
        variational_distribution = CholeskyVariationalDistribution(train_x.size(0))
        variational_strategy = VariationalStrategy(self, train_x, variational_distribution, learn_inducing_locations=False)
        super(GPClassificationModel, 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)
        latent_pred = gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
        return latent_pred
model = GPClassificationModel(X_train)
likelihood = gpytorch.likelihoods.BernoulliLikelihood()

training_iterations = 50
model.train()
likelihood.train()
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
mll = gpytorch.mlls.VariationalELBO(likelihood, model, y_train.numel())
for i in range(training_iterations):
    optimizer.zero_grad()
    output = model(X_train)
    loss = -mll(output, y_train)
    loss.backward()
    print('Iter %d/%d - Loss: %.3f' % (i + 1, training_iterations, loss.item()))
    optimizer.step()

Iter 1/50 - Loss: 0.908
Iter 2/50 - Loss: 1.220
Iter 3/50 - Loss: 0.776
Iter 4/50 - Loss: 0.785
Iter 5/50 - Loss: 0.843
Iter 6/50 - Loss: 0.717
Iter 7/50 - Loss: 0.554
Iter 8/50 - Loss: 0.497
Iter 9/50 - Loss: 0.526
Iter 10/50 - Loss: 0.522
Iter 11/50 - Loss: 0.442
Iter 12/50 - Loss: 0.343
Iter 13/50 - Loss: 0.284
Iter 14/50 - Loss: 0.265
Iter 15/50 - Loss: 0.255
Iter 16/50 - Loss: 0.231
Iter 17/50 - Loss: 0.195
Iter 18/50 - Loss: 0.162
Iter 19/50 - Loss: 0.146
Iter 20/50 - Loss: 0.146
etc.

Loss seems to decrease nicely!!
However this is only on a small subset of the data.
Lets utilize the GPU and train the whole dataset in mini-batches:

Training on GPU

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
device(type='cuda')

X_train.shape, y_train.shape, X_test.shape, y_test.shape
(torch.Size([60912, 240]), torch.Size([60912]), torch.Size([15228, 240]), torch.Size([15228]))

batch_size = 16
from torch.utils.data import TensorDataset, DataLoader
train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_dataset = TensorDataset(X_test, y_test)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
from gpytorch.models import ApproximateGP
from gpytorch.variational import CholeskyVariationalDistribution
from gpytorch.variational import VariationalStrategy
class GPClassificationModel(ApproximateGP):
    def __init__(self, train_x):
        variational_distribution = CholeskyVariationalDistribution(train_x.size(0))
        variational_strategy = VariationalStrategy(self, train_x, variational_distribution, learn_inducing_locations=False)
        super(GPClassificationModel, 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)
        latent_pred = gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
        return latent_pred
model = GPClassificationModel(X_train[:1])
likelihood = gpytorch.likelihoods.BernoulliLikelihood()
model = model.cuda()
likelihood = likelihood.cuda()

epochs = 5
model.train()
likelihood.train()
optimizer = torch.optim.Adam(model.parameters(), lr=0.1)
mll = gpytorch.mlls.VariationalELBO(likelihood, model, y_train.numel())
for epoch in range(epochs):
    running_loss = 0.0
    for i, data in enumerate(train_loader):        
        x_batch, y_batch = data
        x_batch, y_batch = x_batch.to(device), y_batch.to(device)
        optimizer.zero_grad()
        output = model(x_batch)
        loss = -mll(output, y_batch)
        loss.backward()        
        optimizer.step()
        running_loss += loss.item()
        if i % 1000 == 999:    # print every 2000 mini-batches
            print('[%d, %5d] loss: %.3f' % (epoch, i, running_loss / 1000))
            running_loss = 0.0

[0, 999] loss: 0.699
[0, 1999] loss: 0.698
[0, 2999] loss: 0.698
[1, 999] loss: 0.698
[1, 1999] loss: 0.698
[1, 2999] loss: 0.698
[2, 999] loss: 0.698
[2, 1999] loss: 0.699
[2, 2999] loss: 0.698
[3, 999] loss: 0.697
[3, 1999] loss: 0.698
[3, 2999] loss: 0.697
[4, 999] loss: 0.697
[4, 1999] loss: 0.699
[4, 2999] loss: 0.697

Loss doesn't seem to decrease at all!
Now possibly I've got some dumb setting or initialization wrong, I need help here.

Does anyone please have an idea? Im grateful for any comment :) Thanks

[gpleiss]

Try a larger batch size than 16.

[MJimitater]

Hi @gpleiss thanks for your fast reply!!
batch size of 512 started working, then batch size of 1024 showed the hoped-for results!!
Thank you Geoff

BTW how did you know that the problem was in the batch size?

EDIT: I also increased epochs from 5 to 30 in order to see some decent loss decrease

[MJimitater]

Also I realized that initializing with

model = GPClassificationModel(X_train[:1])

is not a good choice, but it should rather be

model = GPClassificationModel(X_train[:batch_size])

to have stable training. Why is this the case?

[gpleiss]

Based on your constructor, the first argument is the set of initial inducing points. X_train[:1] means you'll have a single inducing point, whereas X_train[:batch_size] means you'll have batch_size inducing points.

[MJimitater]

Thanks for your quick reply! As Im still new to the field, I didn't understand the concept of inducing points, a X_train[:1] is simply not reasonable. I will careful pick the inducing points and increase the amount in order to further gain higher accuracies. Thanks again for your help, I appreciate it!