Updated the VAE example with some enhancements

[Coderx7]

Add some enhancements to the VAE example including some comments, a new way for calculating loss (i.e using mean instead of sum, and using mse instead of BCE) and finally several routines for creating visualizations concerning the VAE as other frameworks such as Keras, etc provide.
Here is an example output :

Train Epoch: 50 [0/60000 (0%)]  Loss: 142.077072
Train Epoch: 50 [1280/60000 (2%)]       Loss: 151.412460
Train Epoch: 50 [2560/60000 (4%)]       Loss: 143.169495
Train Epoch: 50 [3840/60000 (6%)]       Loss: 139.854614
Train Epoch: 50 [5120/60000 (9%)]       Loss: 140.479279
Train Epoch: 50 [6400/60000 (11%)]      Loss: 141.113678
Train Epoch: 50 [7680/60000 (13%)]      Loss: 145.222717
Train Epoch: 50 [8960/60000 (15%)]      Loss: 143.298386
Train Epoch: 50 [10240/60000 (17%)]     Loss: 142.074677
Train Epoch: 50 [11520/60000 (19%)]     Loss: 142.046265
Train Epoch: 50 [12800/60000 (21%)]     Loss: 140.014236
Train Epoch: 50 [14080/60000 (23%)]     Loss: 144.403824
Train Epoch: 50 [15360/60000 (26%)]     Loss: 141.879059
Train Epoch: 50 [16640/60000 (28%)]     Loss: 139.471130
Train Epoch: 50 [17920/60000 (30%)]     Loss: 143.647278
Train Epoch: 50 [19200/60000 (32%)]     Loss: 153.391830
Train Epoch: 50 [20480/60000 (34%)]     Loss: 142.550720
Train Epoch: 50 [21760/60000 (36%)]     Loss: 147.450180
Train Epoch: 50 [23040/60000 (38%)]     Loss: 149.538467
Train Epoch: 50 [24320/60000 (41%)]     Loss: 143.521545
Train Epoch: 50 [25600/60000 (43%)]     Loss: 142.479950
Train Epoch: 50 [26880/60000 (45%)]     Loss: 150.655380
Train Epoch: 50 [28160/60000 (47%)]     Loss: 142.734924
Train Epoch: 50 [29440/60000 (49%)]     Loss: 145.209045
Train Epoch: 50 [30720/60000 (51%)]     Loss: 148.113403
Train Epoch: 50 [32000/60000 (53%)]     Loss: 150.475647
Train Epoch: 50 [33280/60000 (55%)]     Loss: 145.669861
Train Epoch: 50 [34560/60000 (58%)]     Loss: 147.789429
Train Epoch: 50 [35840/60000 (60%)]     Loss: 149.124664
Train Epoch: 50 [37120/60000 (62%)]     Loss: 141.129578
Train Epoch: 50 [38400/60000 (64%)]     Loss: 145.054382
Train Epoch: 50 [39680/60000 (66%)]     Loss: 145.409058
Train Epoch: 50 [40960/60000 (68%)]     Loss: 142.284454
Train Epoch: 50 [42240/60000 (70%)]     Loss: 148.351013
Train Epoch: 50 [43520/60000 (72%)]     Loss: 143.500214
Train Epoch: 50 [44800/60000 (75%)]     Loss: 151.315079
Train Epoch: 50 [46080/60000 (77%)]     Loss: 145.327087
Train Epoch: 50 [47360/60000 (79%)]     Loss: 142.971786
Train Epoch: 50 [48640/60000 (81%)]     Loss: 140.635880
Train Epoch: 50 [49920/60000 (83%)]     Loss: 145.872925
Train Epoch: 50 [51200/60000 (85%)]     Loss: 134.699051
Train Epoch: 50 [52480/60000 (87%)]     Loss: 146.803940
Train Epoch: 50 [53760/60000 (90%)]     Loss: 146.638092
Train Epoch: 50 [55040/60000 (92%)]     Loss: 132.822876
Train Epoch: 50 [56320/60000 (94%)]     Loss: 139.588501
Train Epoch: 50 [57600/60000 (96%)]     Loss: 139.583817
Train Epoch: 50 [58880/60000 (98%)]     Loss: 149.037277
====> Epoch: 50 Average loss: 1.1320
====> Test set loss: 148.8819

latent space:

vae animation :

2d digits manifold

Acknowledgement :
Calculating loss using mean (and mse) were inspired by Keras version of VAE
Visualization related parameters and some snippets were inspired (and used) from this blog

Note:
Tested successfully on both Win10 and on Linux (using Google Colab) using Pytorch 1.2.0

[Johnson-yue]

I think the recon loss should not be MSE loss , do you test any dataset except MNIST by using MSE loss and KLD loss?

I experiment result is that the original recon loss (BCE) performance is better than MSE loss on other larger dataset

[Coderx7]

Hi, Yes BCE usually performs better it seems.
I included this because it was also in the keras/tensorflow official repository, I just added it for the sake of completeness and so that users can experiment freely with both options at their disposal.

[Coderx7]

@Johnson-yue : Whats the status of this PR?
Does it not add any value to the current sample we have?
would appreciate if you could decide.
Thanks in advance