Balance VAE losses, add reconstruction + sampling

Author: Kaixhin

vae/.gitignore

@@ -0,0 +1 @@
+*.png

vae/main.py

@@ -2,12 +2,13 @@
 import argparse
 import torch
 import torch.utils.data
-import torch.nn as nn
-import torch.optim as optim
+from torch import nn, optim
 from torch.autograd import Variable
 from torchvision import datasets, transforms
+from torchvision.utils import save_image
 
-parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
+
+parser = argparse.ArgumentParser(description='VAE MNIST Example')
 parser.add_argument('--batch-size', type=int, default=128, metavar='N',
  help='input batch size for training (default: 128)')
 parser.add_argument('--epochs', type=int, default=10, metavar='N',
@@ -54,18 +55,21 @@ def encode(self, x):
  h1 = self.relu(self.fc1(x))
  return self.fc21(h1), self.fc22(h1)
 
- def reparametrize(self, mu, logvar):
- std = logvar.mul(0.5).exp_()
- eps = Variable(std.data.new(std.size()).normal_())
- return eps.mul(std).add_(mu)
+ def reparameterize(self, mu, logvar):
+ if self.training:
+ std = logvar.mul(0.5).exp_()
+ eps = Variable(std.data.new(std.size()).normal_())
+ return eps.mul(std).add_(mu)
+ else:
+ return mu
 
  def decode(self, z):
  h3 = self.relu(self.fc3(z))
  return self.sigmoid(self.fc4(h3))
 
  def forward(self, x):
  mu, logvar = self.encode(x.view(-1, 784))
- z = self.reparametrize(mu, logvar)
+ z = self.reparameterize(mu, logvar)
  return self.decode(z), mu, logvar
 
 
@@ -74,7 +78,6 @@ def forward(self, x):
  model.cuda()
 
 reconstruction_function = nn.BCELoss()
-reconstruction_function.size_average = False
 
 
 def loss_function(recon_x, x, mu, logvar):
@@ -86,6 +89,8 @@ def loss_function(recon_x, x, mu, logvar):
  # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2)
  KLD_element = mu.pow(2).add_(logvar.exp()).mul_(-1).add_(1).add_(logvar)
  KLD = torch.sum(KLD_element).mul_(-0.5)
+ # Normalise by same number of elements as in reconstruction
+ KLD /= args.batch_size * 784
 
  return BCE + KLD
 
@@ -119,12 +124,15 @@ def train(epoch):
 def test(epoch):
  model.eval()
  test_loss = 0
- for data, _ in test_loader:
+ for i, (data, _) in enumerate(test_loader):
  if args.cuda:
  data = data.cuda()
  data = Variable(data, volatile=True)
  recon_batch, mu, logvar = model(data)
  test_loss += loss_function(recon_batch, data, mu, logvar).data[0]
+ if i == 0:
+ save_image(recon_batch.data.view(args.batch_size, 1, 28, 28),
+ 'reconstruction_' + str(epoch) + '.png')
 
  test_loss /= len(test_loader.dataset)
  print('====> Test set loss: {:.4f}'.format(test_loss))
@@ -133,3 +141,5 @@ def test(epoch):
 for epoch in range(1, args.epochs + 1):
  train(epoch)
  test(epoch)
+ sample = model.decode(Variable(torch.randn(64, 20)))
+ save_image(sample.data.view(64, 1, 28, 28), 'sample_' + str(epoch) + '.png')