| 116 | |
| 117 | # # Reconstruction + KL divergence losses summed over all elements and batch |
| 118 | def loss_function(recon_x, x, mu, logvar): |
| 119 | BCE = dy.binary_log_loss(recon_x, x) # equiv to torch.nn.functional.binary_cross_entropy(?,?, size_average=False) |
| 120 | # see Appendix B from VAE paper: |
| 121 | # Kingma and Welling. Auto-Encoding Variational Bayes. ICLR, 2014 |
| 122 | # https://arxiv.org/abs/1312.6114 |
| 123 | # 0.5 * sum(1 + log(sigma^2) - mu^2 - sigma^2) |
| 124 | KLD = -0.5 * dy.sum_elems(1 + logvar - dy.pow(mu, dy.scalarInput(2)) - dy.exp(logvar)) |
| 125 | |
| 126 | return BCE + KLD |
| 127 | |
| 128 | |
| 129 | def train(epoch): |