GPytorch + Pytorch Lightening

[andreimargeloiu]

Are there known issues of using Gpytorch with Pytorch lightning?

For example, is there any functionality of a gpytorch.Module that wouldn't work if it's defined as an attribute of a class extending LighteningModule?

[gpleiss]

I'm not sure if anyone has attempted a Gpytorch/lightning integration yet. I'm curious to know if it does work or if there are any issues that we'd need to fix.

[andreimargeloiu]

I am using Stochastic Deep Kernel learning with PytorchLightening and haven't encountered any problem. Specifically, the model class inherits from two classes:

class DKLModelLightening(gpytorch.Module, pl.LightningModule)

[nilsleh]

@andreimargeloiu I am also trying to implement Stochastic Deep Kernel Learning with a lightning module, but struggling to get it to work. I made a google colab notebook that tries to implement it on a Toy Regression Dataset but I just get mean 0 predictions. I am not sure what I am doing wrong, and was wondering whether you had any pointers? Thanks in advance.

[jacobrgardner]

@nilsleh I might have missed it, but I don't see where the gpytorch model is being put into eval mode? In train mode, you'll get the prior when you call the model, which could definitely have a zero mean everywhere.

[nilsleh]

Thanks you for your reply @jacobrgardner . My first assumption was that eval and train mode is something that is being handled by pytorch lightning as one usually does not have to do this explicitly. However, also when setting eval mode explicitly (updated notebook) the results do not change. So I am suspecting that I am either setting up the gpytorch approximate GP incorrectly or that there is something with lightning happening under the hood that I am not aware of.

[andreimargeloiu]

I used the following code (with pytorch-lightning 1.5.9), and I didn't have any issues:

class GaussianProcessLayer(gpytorch.models.ApproximateGP):
	def __init__(self, args):
		variational_distribution = gpytorch.variational.CholeskyVariationalDistribution(
			num_inducing_points=args.grid_size,
			batch_shape=torch.Size([args.num_tasks])
		)

		# 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=args.grid_size,
				grid_bounds=[(-args.grid_bound, args.grid_bound)],
				variational_distribution=variational_distribution,
			), num_tasks=args.num_tasks,
		)
		super().__init__(variational_strategy)

		self.mean_module = gpytorch.means.ConstantMean()
		self.covar_module = gpytorch.kernels.ScaleKernel(
			gpytorch.kernels.RBFKernel(	# TODO add ARD on each dimension
				lengthscale_prior=gpytorch.priors.SmoothedBoxPrior(
					math.exp(-1), math.exp(1), sigma=0.1, transform=torch.exp
				)
			)
		)
		
	def forward(self, x):
		mean = self.mean_module(x)
		covar = self.covar_module(x)
		return gpytorch.distributions.MultivariateNormal(mean, covar)


class DKLModel(gpytorch.Module, pl.LightningModule):
	def __init__(self, args, feature_extractor, decoder):
		super().__init__()

		self.args = args
		self.feature_extractor = feature_extractor
		self.decoder = decoder

		self.gp_layer = GaussianProcessLayer(args)
		self.likelihood = gpytorch.likelihoods.SoftmaxLikelihood(  	# The softmax likelihood includes an additional mixing layer linear layer 
			num_features=args.num_tasks,
			num_classes=args.num_classes
		)
		self.mll_elbo = gpytorch.mlls.VariationalELBO(
				self.likelihood, 
				self.gp_layer,
				num_data=args.train_size,
				combine_terms=False
		)

		self.learning_rate = args.lr
		self.steps = args.max_steps
		self.epochs = args.max_epochs
		self.log_test_key = None

		# This module will scale the NN features so that they're nice values
		self.scale_to_bounds = gpytorch.utils.grid.ScaleToBounds(-args.grid_bound, args.grid_bound)
	
	def forward(self, x):
		hidden_rep, features_extracted = self.feature_extractor(x)
		
		x_hat = self.decoder(hidden_rep) if self.decoder else None      			# reconstruction

		features_extracted = self.scale_to_bounds(features_extracted)				# for stability
		features_extracted = features_extracted.transpose(-1, -2).unsqueeze(-1)		# Makes it so that we learn a GP for each feature
		gp_output = self.gp_layer(features_extracted)
		return gp_output, x_hat

	def compute_loss(self, y_true, gp_output, x, x_hat):
		log_likelihood, kl_divergence, log_prior = self.mll_elbo(gp_output, y_true)
		
		mll = log_likelihood - kl_divergence + log_prior
		reconstruction_loss = F.mse_loss(x_hat, x, reduction='mean') if self.decoder else torch.zeros(1, device=self.device)

		return {
			'nll': -log_likelihood,
			'kl': kl_divergence,
			'mll': mll,
			'reconstruction': self.args.gamma * reconstruction_loss,
			'total': -mll + self.args.gamma * reconstruction_loss
		}

	def log_losses(self, losses, key):
		self.log(f"{key}/total_loss", losses['total'])
		self.log(f"{key}/reconstruction_loss", losses['reconstruction'])
		self.log(f"{key}/marginal_likelihood", losses['mll'])
		self.log(f"{key}/nll", losses['nll'])
		self.log(f"{key}/kl", losses['kl'])

	def log_epoch_metrics(self, outputs, key):
		y_true, y_pred = get_labels_lists(outputs)

		metrics = compute_all_metrics(self.args, y_true, y_pred)
		for metric_name, metric_value in metrics.items():
			self.log(f"{key}/{metric_name}", metric_value)
			

	def training_step(self, batch, batch_idx):
		x, y_true = batch
		with gpytorch.settings.num_likelihood_samples(48), gpytorch.settings.cholesky_jitter(1e-5):
			gp_output, x_hat = self.forward(x) 				        # MultivariateNormal
			losses = self.compute_loss(y_true, gp_output, x, x_hat)

			self.log_losses(losses, 'train')

			with torch.no_grad():
				output_samples = self.likelihood(gp_output) 		# draw samples from the posterior
				y_pred = output_samples.probs.mean(0).argmax(-1)  	# Taking the mean over all of the sample we've drawn

			return {
				'loss': losses['total'],
				'losses': detach_tensors(losses),
				'y_true': y_true,
				'y_pred': y_pred
			}

	def training_epoch_end(self, outputs):
		self.log_epoch_metrics(outputs, 'train')

	def validation_step(self, batch, batch_idx):
		x, y_true = batch
		with gpytorch.settings.num_likelihood_samples(48), gpytorch.settings.cholesky_jitter(1e-5):
			gp_output, x_hat = self.forward(x) 						# MultivariateNormal
			losses = self.compute_loss(y_true, gp_output, x, x_hat)
			
			with torch.no_grad():
				output_samples = self.likelihood(gp_output) 		# draw samples from the posterior
				y_pred = output_samples.probs.mean(0).argmax(-1)    # Taking the mean over all of the sample we've drawn

			return {
				'losses': detach_tensors(losses),
				'y_true': y_true,
				'y_pred': y_pred
			}

	def validation_epoch_end(self, outputs):
		losses = {
			'total': np.mean([output['losses']['total'].item() for output in outputs]),
			'mll': np.mean([output['losses']['mll'].item() for output in outputs]),
			'nll': np.mean([output['losses']['nll'].item() for output in outputs]),
			'kl': np.mean([output['losses']['kl'].item() for output in outputs]),
			'reconstruction': np.mean([output['losses']['reconstruction'].item() for output in outputs]),
		}

		self.log_losses(losses, key='valid')
		self.log_epoch_metrics(outputs, 'valid')

	def test_step(self, batch, batch_idx):
		x, y_true = batch
		with gpytorch.settings.num_likelihood_samples(48):
			gp_output, x_hat = self.forward(x) 				# MultivariateNormal

			with torch.no_grad():
				output_samples = self.likelihood(gp_output) # draw samples from the posterior
				y_pred = output_samples.probs.mean(0).argmax(-1)  # Taking the mean over all of the sample we've drawn

			return {
				'y_true': y_true,
				'y_pred': y_pred
			}

	def test_epoch_end(self, outputs):
		self.log_epoch_metrics(outputs, self.log_test_key)

	def configure_optimizers(self):
		optimizer = Adam([
			{'params': self.feature_extractor.parameters(), 'weight_decay': 1e-4},             # feature extractor
			{'params': self.gp_layer.hyperparameters(), 'lr': self.learning_rate * 0.01},      # GP kernel hyperparameters
			{'params': self.gp_layer.variational_parameters()},                                # variational mean, cholesky variational covariance
			{'params': self.likelihood.parameters()},                                        # linear layer from GP output to classes
		], lr=self.learning_rate, weight_decay=0)
		
		# lr_scheduler = MultiStepLR(optimizer, milestones=[0.5 * self.epochs, 0.75 * self.epochs], gamma=0.1)

		if self.args.lookahead_optimizer:
			optimizer = Lookahead(optimizer, la_steps=5, la_alpha=0.5)

		return {
			'optimizer': optimizer,
			# 'lr_scheduler': lr_scheduler
		}

[nilsleh]

Thanks a lot for sending your code snippet @andreimargeloiu , I want to say that I am using the same code structure in my example as you showed. However, predictions are still crappy, so I think I am setting up the DKL Approximate GP in a wrong manner. The predictions are not strictly the prior output, but it does not learn the linear trend. Do you happen to have a DKL with Approximate GP regression example @jacobrgardner (not with lightning just gpytorch) that could potentially narrow down my issue?

[nilsleh]

It turns out the implementation was working but the Deep Kernel Learning model is very sensitive to hyperparameters, especially the batch size and the learning rate. As to why this is, would probably be a better topic for another discussion and is very much a question I am wondering about. However, after some exploring there exists hyperparameters that yield a desireable result, which can be seen in the updated notebook version.

[nm19000]

Hi Andrei, thank you for your answer and code example. May I ask why you set ```combine_terms=False````? in

self.mll_elbo = gpytorch.mlls.VariationalELBO(
self.likelihood,
self.gp_layer,
num_data=args.train_size,
combine_terms=False
)
?