fix test marginal likelihood

Author: Jaan Altosaar

train_variational_autoencoder_pytorch.py

@@ -1,17 +1,19 @@
-"""Fit a VAE to MNIST. 
+"""Fit a variational autoencoder to MNIST. 
 
-Conventions:
+Notes:
+  - run https://github.com/altosaar/proximity_vi/blob/master/get_binary_mnist.py to download binary MNIST file
   - batch size is the innermost dimension, then the sample dimension, then latent dimension
 """
 import torch
 import torch.utils
+import torch.utils.data
 from torch import nn
 import nomen
 import yaml
 import numpy as np
 import logging
-
-import data
+import pathlib
+import h5py
 
 config = """
 latent_size: 128
@@ -20,96 +22,118 @@
 batch_size: 128
 test_batch_size: 512
 max_iterations: 100000
-log_interval: 1000
-n_samples: 77
+log_interval: 5000
+n_samples: 128
+use_gpu: true
+train_dir: $TMPDIR
 """
 
-class NeuralNetwork(nn.Module):
-  def __init__(self, input_size, output_size, hidden_size):
-    super().__init__()
-    modules = [nn.Linear(input_size, hidden_size),
-               nn.ReLU(),
-               nn.Linear(hidden_size, hidden_size),
-               nn.ReLU(),
-               nn.Linear(hidden_size, output_size)]
-    self.net = nn.Sequential(*modules)
-    
-  def forward(self, input):
-    return self.net(input)
-
-
 
 class Model(nn.Module):
   """Bernoulli model parameterized by a generative network with Gaussian latents for MNIST."""
-  def __init__(self, latent_size, data_size, batch_size):
+  def __init__(self, latent_size, data_size, batch_size, device):
     super().__init__()
-    # prior on latents is standard normal
-    self.p_z = torch.distributions.Normal(torch.zeros(latent_size), torch.ones(latent_size))
-    # likelihood is bernoulli, equivalent to negative binary cross entropy
+    self.p_z = torch.distributions.Normal(
+        torch.zeros(latent_size, device=device), 
+        torch.ones(latent_size, device=device))
     self.log_p_x = BernoulliLogProb()
-    # generative network is a MLP
-    self.generative_network = NeuralNetwork(input_size=latent_size, output_size=data_size, hidden_size=latent_size * 2)
-    
+    self.generative_network = NeuralNetwork(input_size=latent_size,
+                                            output_size=data_size, 
+                                            hidden_size=latent_size * 2)
 
   def forward(self, z, x):
     """Return log probability of model."""
     log_p_z = self.p_z.log_prob(z).sum(-1)
     logits = self.generative_network(z)
+    # unsqueeze sample dimension
+    logits, x = torch.broadcast_tensors(logits, x.unsqueeze(1))
     log_p_x = self.log_p_x(logits, x).sum(-1)
     return log_p_z + log_p_x
 
-
-class NormalLogProb(nn.Module):
-  def __init__(self):
-    super().__init__()
-
-  def forward(self, loc, scale, z):
-    var = torch.pow(scale, 2)
-    return -0.5 * torch.log(2 * np.pi * var) + torch.pow(z - loc, 2) / (2 * var)
-
-class BernoulliLogProb(nn.Module):
-  def __init__(self):
-    super().__init__()
-    self.bce_with_logits = nn.BCEWithLogitsLoss(reduction='none')
-
-  def forward(self, logits, target):
-    logits, target = torch.broadcast_tensors(logits, target.unsqueeze(1))
-    return -self.bce_with_logits(logits, target)
 
 class Variational(nn.Module):
   """Approximate posterior parameterized by an inference network."""
   def __init__(self, latent_size, data_size):
     super().__init__()
-    self.inference_network = NeuralNetwork(input_size=data_size, output_size=latent_size * 2, hidden_size=latent_size*2)
+    self.inference_network = NeuralNetwork(input_size=data_size, 
+                                           output_size=latent_size * 2, 
+                                           hidden_size=latent_size*2)
     self.log_q_z = NormalLogProb()
     self.softplus = nn.Softplus()
 
   def forward(self, x, n_samples=1):
     """Return sample of latent variable and log prob."""
     loc, scale_arg = torch.chunk(self.inference_network(x).unsqueeze(1), chunks=2, dim=-1)
     scale = self.softplus(scale_arg)
-    eps = torch.randn((loc.shape[0], n_samples, loc.shape[-1]))
+    eps = torch.randn((loc.shape[0], n_samples, loc.shape[-1]), device=loc.device)
     z = loc + scale * eps  # reparameterization
     log_q_z = self.log_q_z(loc, scale, z).sum(-1)
     return z, log_q_z
 
 
+class NeuralNetwork(nn.Module):
+  def __init__(self, input_size, output_size, hidden_size):
+    super().__init__()
+    modules = [nn.Linear(input_size, hidden_size),
+               nn.ReLU(),
+               nn.Linear(hidden_size, hidden_size),
+               nn.ReLU(),
+               nn.Linear(hidden_size, output_size)]
+    self.net = nn.Sequential(*modules)
+    
+  def forward(self, input):
+    return self.net(input)
+
+
+class NormalLogProb(nn.Module):
+  def __init__(self):
+    super().__init__()
+
+  def forward(self, loc, scale, z):
+    var = torch.pow(scale, 2)
+    return -0.5 * torch.log(2 * np.pi * var) - torch.pow(z - loc, 2) / (2 * var)
+
+
+class BernoulliLogProb(nn.Module):
+  def __init__(self):
+    super().__init__()
+    self.bce_with_logits = nn.BCEWithLogitsLoss(reduction='none')
+
+  def forward(self, logits, target):
+    # bernoulli log prob is equivalent to negative binary cross entropy
+    return -self.bce_with_logits(logits, target)
+
+
 def cycle(iterable):
   while True:
     for x in iterable:
       yield x
 
 
+def load_binary_mnist(cfg, **kwcfg):
+  f = h5py.File(pathlib.os.path.join(pathlib.os.environ['DAT'], 'binarized_mnist.hdf5'), 'r')
+  x_train = f['train'][::]
+  x_val = f['valid'][::]
+  x_test = f['test'][::]
+  train = torch.utils.data.TensorDataset(torch.from_numpy(x_train))
+  train_loader = torch.utils.data.DataLoader(train, batch_size=cfg.batch_size, shuffle=True)
+  validation = torch.utils.data.TensorDataset(torch.from_numpy(x_val))
+  val_loader = torch.utils.data.DataLoader(validation, batch_size=cfg.test_batch_size, shuffle=False)
+  test = torch.utils.data.TensorDataset(torch.from_numpy(x_test))
+  test_loader = torch.utils.data.DataLoader(test, batch_size=cfg.test_batch_size, shuffle=False)
+  return train_loader, val_loader, test_loader
+
+
 def evaluate(n_samples, model, variational, eval_data):
   model.eval()
   total_log_p_x = 0.0
   total_elbo = 0.0
   for batch in eval_data:
-    x = batch[0]
+    x = batch[0].to(next(model.parameters()).device)
     z, log_q_z = variational(x, n_samples)
     log_p_x_and_z = model(z, x)
-    # importance sampling of approximate marginal likelihood
-    # using logsumexp in the sample dimension
+    # importance sampling of approximate marginal likelihood with q(z)
+    # as the proposal, and logsumexp in the sample dimension
     elbo = log_p_x_and_z - log_q_z
     log_p_x = torch.logsumexp(elbo, dim=1) - np.log(n_samples)
     # average over sample dimension, sum over minibatch
@@ -123,28 +147,59 @@ def evaluate(n_samples, model, variational, eval_data):
 if __name__ == '__main__':
   dictionary = yaml.load(config)
   cfg = nomen.Config(dictionary)
-  
-  model = Model(latent_size=cfg.latent_size, data_size=cfg.data_size, batch_size=cfg.batch_size)
-  variational = Variational(latent_size=cfg.latent_size, data_size=cfg.data_size)
+  device = torch.device("cuda:0" if cfg.use_gpu else "cpu")
+
+  model = Model(latent_size=cfg.latent_size, 
+                data_size=cfg.data_size,
+                batch_size=cfg.batch_size, 
+                device=device)
+  variational = Variational(latent_size=cfg.latent_size,
+                            data_size=cfg.data_size)
+  model.to(device)
+  variational.to(device)
+
+  optimizer = torch.optim.RMSprop(list(model.parameters()) +
+                                  list(variational.parameters()),
+                                  lr=cfg.learning_rate,
+                                  centered=True)
 
-  optimizer = torch.optim.RMSprop(list(model.parameters()) + list(variational.parameters()), 
-                                  lr=cfg.learning_rate)
+  kwargs = {'num_workers': 0, 'pin_memory': False} if cfg.use_gpu else {}
+  train_data, valid_data, test_data = load_binary_mnist(cfg, **kwargs)
 
-  train_data, valid_data, test_data = data.load_binary_mnist(cfg)
+  best_valid_elbo = -np.inf
+  num_no_improvement = 0
 
   for step, batch in enumerate(cycle(train_data)):
-    x = batch[0]
+    x = batch[0].to(device)
     model.zero_grad()
     variational.zero_grad()
     z, log_q_z = variational(x)
     log_p_x_and_z = model(z, x)
+    # average over sample dimension
     elbo = (log_p_x_and_z - log_q_z).mean(1)
-    loss = -elbo.mean(0)
+    # sum over batch dimension
+    loss = -elbo.sum(0)
     loss.backward()
     optimizer.step()
 
     if step % cfg.log_interval == 0:
-      print(f'step:\t{step}\ttrain elbo: {elbo.detach().cpu().numpy()[0]:.2f}')
+      print(f'step:\t{step}\ttrain elbo: {elbo.detach().cpu().numpy().mean():.2f}')
       with torch.no_grad():
         valid_elbo, valid_log_p_x = evaluate(cfg.n_samples, model, variational, valid_data)
-      print(f'step:\t{step}\tvalid elbo: {valid_elbo:.2f}\tvalid log p(x): {valid_log_p_x:.2f}')
+      print(f'step:\t{step}\t\tvalid elbo: {valid_elbo:.2f}\tvalid log p(x): {valid_log_p_x:.2f}')
+      if valid_elbo > best_valid_elbo:
+        best_valid_elbo = valid_elbo
+        states = {'model': model.state_dict(),
+                  'variational': variational.state_dict()}
+        torch.save(states, cfg.train_dir / 'best_state_dict')
+      else:
+        num_no_improvement += 1
+
+      if num_no_improvement > 5:
+        checkpoint = torch.load(cfg.train_dir / 'best_state_dict')
+        model.load_state_dict(checkpoint['model'])
+        variational.load_state_dict(checkpoint['variational'])
+        with torch.no_grad():
+          test_elbo, test_log_p_x = evaluate(cfg.n_samples, model, variational, test_data)
+        print(f'step:\t{step}\t\ttest elbo: {test_elbo:.2f}\ttest log p(x): {test_log_p_x:.2f}')
+        break