def train_model(self, dataset, epoch, optimizer, visualize_errors=False):
"""
Trains a given model for a given number of epochs on a given dataset.
"""
# self.train()
super().train()
dataset_size = dataset.features.shape[0]
# batch_size = dataset_size # train on entire data in one batch
batch_size = self.hyperparams.batch_size # train using the specified batch size
# transform dataset to dataloader
train_loader = dataset_to_dataloader(dataset, batch_size=batch_size)
n_train_samples = len(train_loader.dataset)
n_batches = len(train_loader)
mse = 0
mae = 0
total_loss = 0
total_log_likelihood = 0
total_kl_divergence = 0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(self.device), target.to(self.device)
self.zero_grad()
loss, log_prior, log_variational_posterior, log_likelihood, outputs = self.sample_elbo(data, target, dataset_size)
loss.backward()
optimizer.step()
# perform analytic update of remaining model parameters
self.analytic_update()
# Given all model outputs, compute the mean output of the ensemble
mean_output = outputs.mean(dim=1)
device_type = self.device.type
mse, mae = update_mse_mae(mse, mae, device_type, mean_output, target)
total_loss += loss
total_log_likelihood += -log_likelihood
total_kl_divergence += (log_variational_posterior - log_prior) * target.size()[0] / dataset_size
rmse = np.sqrt(mse / dataset_size)
mae /= dataset_size
if visualize_errors:
self.train_writer.add_scalar('loss__training loss', total_loss.item(), epoch)
self.train_writer.add_scalar('loss__kl term' , total_kl_divergence.item(), epoch)
self.train_writer.add_scalar('loss__log_likelihood term', total_log_likelihood.item(), epoch)
if not self.hyperparams.classification:
self.train_writer.add_scalar('errors__mae', mae, epoch)
self.train_writer.add_scalar('errors__rmse', rmse, epoch)
return loss, rmse, mae
def predict(self, dataset, epoch=1, mean_y_train=0, std_y_train=1, visualize_errors=False):
"""
Evaluates an ensemble of networks on a given test dataset.
Because the Bayesian Neural Network has a distribution over the weights,
we basically have an infinite number of different neural networks. We can
take advantage of that by using an ensemble of networks during prediction.
Each model in the ensemble performs a prediction. The different predictions
are then averaged to give a final output.
A model can be obtained by sampling weights from the distribution. Note: for
each input batch, new models will be sampled.
"""
n_samples_testing = self.hyperparams.n_samples_testing
dataset_size = dataset.features.shape[0]
test_batch_size = dataset_size
# transform dataset to dataloader
test_loader = dataset_to_dataloader(dataset, batch_size=test_batch_size, shuffle=False)
n_test_samples = len(test_loader.dataset)
n_test_batches = len(test_loader)
super(GaussianBNN, self).eval()
mse = 0
rmse = 0
mae = 0
loglike = 0
all_predicted_distributions = []
means = []
with torch.no_grad():
for data, target in test_loader:
data, target = data.to(self.device), target.to(self.device)
# Each batch is forwarded through each model in the ensemble
# and the model outputs are saved.
ensemble_outputs = self.forward(data, sample=True, n_samples=n_samples_testing) * std_y_train + mean_y_train
ensemble_outputs = ensemble_outputs.reshape(n_samples_testing, test_batch_size).t()
# calculation of the predictive log likelihood of a batch, see notes from 18.12.18
var_noise = np.exp(self.log_var_noise.detach().numpy()) * std_y_train ** 2
if self.hyperparams.classification:
loglike_factor = - F.binary_cross_entropy_with_logits(ensemble_outputs, target.reshape(-1,1).repeat(1,n_samples_testing), reduction='none')
loglike = torch.sum(torch.logsumexp(loglike_factor - math.log(n_samples_testing), 1))
# Given all model outputs, compute the mean output of the ensemble
mean_output = ensemble_outputs.mean(1)
else:
if self.device.type == 'cuda':
target = target.cpu().numpy()
ensemble_outputs = ensemble_outputs.cpu().numpy()
log_factor = -0.5 * np.log(2 * math.pi * var_noise) - (np.tile(target.reshape(-1, 1), (1, n_samples_testing)) - np.array(ensemble_outputs))**2 / (2* var_noise)
loglike += np.sum(logsumexp(log_factor - np.log(n_samples_testing), 1))
# Given all model outputs, compute the mean output of the ensemble
mean_output = ensemble_outputs.mean(1)
if self.device.type == 'cuda':
target = torch.from_numpy(target)
mean_output = torch.from_numpy(mean_output)
if self.hyperparams.classification:
distributions = [BinarySampleDistribution(1 / (1 + np.exp(-e))) for e in ensemble_outputs.cpu().detach().numpy()]
else:
mse += F.mse_loss(mean_output, target, reduction='sum')
mae += F.l1_loss(mean_output, target, reduction='sum')
distributions = [SampleDistribution(ensemble_outputs[i], var_noise)
for i in range(test_batch_size)]
all_predicted_distributions.extend(distributions)
predicted_distr = PredictiveDistribution(all_predicted_distributions)
loglike /= dataset_size
rmse = np.sqrt(mse / dataset_size)
mae /= dataset_size
if self.hyperparams.classification:
zero_one = AllMetrics.zero_one_loss.compute(target.cpu().detach().numpy(), predicted_distr)
if visualize_errors:
self.test_writer.add_scalar('errors__predictive log likelihood', loglike, epoch)
if self.hyperparams.classification:
self.test_writer.add_scalar('errors__zero_one', zero_one, epoch)
else:
self.test_writer.add_scalar('errors__mae', mae, epoch)
self.test_writer.add_scalar('errors__rmse', rmse, epoch)
return predicted_distr, rmse, mae, -loglike