Fix VAE loss + improve reconstruction viz

Closes pytorch#225

Author: Kaixhin

vae/main.py

@@ -4,6 +4,7 @@
 import torch.utils.data
 from torch import nn, optim
 from torch.autograd import Variable
+from torch.nn import functional as F
 from torchvision import datasets, transforms
 from torchvision.utils import save_image
 
@@ -77,18 +78,15 @@ def forward(self, x):
 if args.cuda:
     model.cuda()
 
-reconstruction_function = nn.BCELoss()
-
 
 def loss_function(recon_x, x, mu, logvar):
-    BCE = reconstruction_function(recon_x, x.view(-1, 784))
+    BCE = F.binary_cross_entropy(recon_x, x.view(-1, 784))
 
     # see Appendix B from VAE paper:
     # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014
     # https://arxiv.org/abs/1312.6114
     # 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)
+    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
     # Normalise by same number of elements as in reconstruction
     KLD /= args.batch_size * 784
 
@@ -131,8 +129,11 @@ def test(epoch):
         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.cpu().view(args.batch_size, 1, 28, 28),
-                     'reconstruction_' + str(epoch) + '.png')
+          n = min(data.size(0), 8)
+          comparison = torch.cat([data[:n],
+                                  recon_batch.view(args.batch_size, 1, 28, 28)[:n]])
+          save_image(comparison.data.cpu(),
+                     'results/reconstruction_' + str(epoch) + '.png', nrow=n)
 
     test_loss /= len(test_loader.dataset)
     print('====> Test set loss: {:.4f}'.format(test_loss))
@@ -145,4 +146,5 @@ def test(epoch):
     if args.cuda:
        sample = sample.cuda()
     sample = model.decode(sample).cpu()
-    save_image(sample.data.view(64, 1, 28, 28), 'sample_' + str(epoch) + '.png')
+    save_image(sample.data.view(64, 1, 28, 28),
+               'results/sample_' + str(epoch) + '.png')