|
| 1 | +# VAE implementation exercise for Depth First Learning Curriculum: Normalizing Flows for Variational Inference. |
| 2 | +# This is a stripped-down version of the PyTorch VAE example available at https://github.com/pytorch/examples/blob/master/vae/main.py. |
| 3 | + |
| 4 | +from __future__ import print_function |
| 5 | +import argparse |
| 6 | +import torch |
| 7 | +import torch.utils.data |
| 8 | +from torch import nn, optim |
| 9 | +from torch.nn import functional as F |
| 10 | +from torchvision import datasets, transforms |
| 11 | +from torchvision.utils import save_image |
| 12 | + |
| 13 | + |
| 14 | +parser = argparse.ArgumentParser(description='VAE MNIST Example') |
| 15 | +parser.add_argument('--batch-size', type=int, default=128, metavar='N', |
| 16 | + help='input batch size for training (default: 128)') |
| 17 | +parser.add_argument('--epochs', type=int, default=10, metavar='N', |
| 18 | + help='number of epochs to train (default: 10)') |
| 19 | +parser.add_argument('--no-cuda', action='store_true', default=False, |
| 20 | + help='enables CUDA training') |
| 21 | +parser.add_argument('--seed', type=int, default=1, metavar='S', |
| 22 | + help='random seed (default: 1)') |
| 23 | +parser.add_argument('--log-interval', type=int, default=10, metavar='N', |
| 24 | + help='how many batches to wait before logging training status') |
| 25 | +args = parser.parse_args() |
| 26 | +args.cuda = not args.no_cuda and torch.cuda.is_available() |
| 27 | + |
| 28 | +torch.manual_seed(args.seed) |
| 29 | + |
| 30 | +device = torch.device("cuda" if args.cuda else "cpu") |
| 31 | + |
| 32 | +kwargs = {'num_workers': 1, 'pin_memory': True} if args.cuda else {} |
| 33 | +train_loader = torch.utils.data.DataLoader( |
| 34 | + datasets.MNIST('../data', train=True, download=True, |
| 35 | + transform=transforms.ToTensor()), |
| 36 | + batch_size=args.batch_size, shuffle=True, **kwargs) |
| 37 | +test_loader = torch.utils.data.DataLoader( |
| 38 | + datasets.MNIST('../data', train=False, transform=transforms.ToTensor()), |
| 39 | + batch_size=args.batch_size, shuffle=True, **kwargs) |
| 40 | + |
| 41 | + |
| 42 | +class VAE(nn.Module): |
| 43 | + def __init__(self): |
| 44 | + super(VAE, self).__init__() |
| 45 | + |
| 46 | + # fc1, fc21 and fc22 are used by the encoder. |
| 47 | + # fc1 takes a vectorized MNIST image as input |
| 48 | + # fc21 and fc22 are both attached to the activation output of fc1 (using ReLU). |
| 49 | + # fc21 outputs the means, and fc22 the log-variances of |
| 50 | + # each component of th 20-dimensional latent Gaussian. |
| 51 | + self.fc1 = nn.Linear(784, 400) |
| 52 | + self.fc21 = nn.Linear(400, 20) |
| 53 | + self.fc22 = nn.Linear(400, 20) |
| 54 | + # fc3 and fc4 are connected in series as the decoder. |
| 55 | + # fc3 takes a realization from the latent space as input |
| 56 | + # and the decoder generates a vectorized 28x28 image. |
| 57 | + # The output of fc3 passes through a ReLU, |
| 58 | + # while fc4 uses a sigmoid in order to output a probability for each pixel |
| 59 | + self.fc3 = nn.Linear(20, 400) |
| 60 | + self.fc4 = nn.Linear(400, 784) |
| 61 | + |
| 62 | + # TODO: Implement the following four functions. Note that they should be able to accept arguments containing stacked information for multiple observations |
| 63 | + # e.g. a minibatch rather than a single observation. Your solution will need to handle this. If you treat the arguments as |
| 64 | + # representing a single observation in your logic, in most cases broadcasting will do the rest of the job automatically for you. |
| 65 | + def encode(self, x): |
| 66 | + # This should return the outputs of fc21 and fc22 as a tuple |
| 67 | + pass |
| 68 | + |
| 69 | + def reparameterize(self, mu, logvar): |
| 70 | + # This should sample vectors from an isotropic Gaussian, and use these to generate |
| 71 | + # and return observations with a mean vectors from mu, and log-variances of log-var |
| 72 | + pass |
| 73 | + |
| 74 | + def decode(self, z): |
| 75 | + # Pass z through the decoder. For each 20-dimensional latent realization, there should be a 784-dimensional vector of |
| 76 | + #probabilities generated, one per pixel |
| 77 | + pass |
| 78 | + |
| 79 | + def forward(self, x): |
| 80 | + # For each observation in x: |
| 81 | + # 1. Pass it through the encoder to get predicted variational distribution parameters |
| 82 | + # 2. Reparameterize an isotropic Gaussian with these parameters to get sample latent variable realizations |
| 83 | + # 3. Pass the realization through the encoder to get predicted pixel probabilities |
| 84 | + # Return a tuple with 3 elements: (a) the predicted pixel probabilities, (b) the predicted variational means, and (c) the predicted variational log-variances |
| 85 | + x = x.view(-1,784) # Reshape x to provide suitable inputs to the encoder |
| 86 | + pass |
| 87 | + |
| 88 | +model = VAE().to(device) |
| 89 | +optimizer = optim.Adam(model.parameters(), lr=1e-3) |
| 90 | + |
| 91 | +# TODO: Implement this loss function |
| 92 | +def loss_function(recon_x, x, mu, logvar): |
| 93 | + # The loss should be (an estimate of) the negative ELBO - remember we wish to maximise the ELBO - but the ELBO can be written in a number of forms. |
| 94 | + # In this case, the prior for the latent variable and the variational posterior are both Gaussians, and we will exploit this. |
| 95 | + # Specifically, we can analytically calculate a part of the ELBO, and only use Monte Carlo estimation for the rest. |
| 96 | + # 1. We use the form of the ELBO which includes a KL divergence between the latent prior and the variational family |
| 97 | + # - see the form at the bottom of page 6 of Blei et al's "Variational Inference: A Review for Statisticians". |
| 98 | + # 2. In this case, the expression for the relevant KL divergence can be obtained from Exercise (e) in Week 1. |
| 99 | + # |
| 100 | + # The other term is the expected conditional log-likelihood, which is estimated using a single Monte-Carlo sample. |
| 101 | + # For the log-likelihood, one evaluates the probability of observing an input point given the "conditional distribution" for |
| 102 | + # observations output by the network - in this case, each pixel is independently Bernoulli with parameter equal to the output probability. |
| 103 | + # You may find torch.nn.functional's binary_cross_entropy function useful here. |
| 104 | + # |
| 105 | + # Additional: the extraction of the KL divergence as above reduces the variance. Investigate the effect of directly estimating |
| 106 | + # the full ELBO term for each observation with a single Monte Carlo sample. |
| 107 | + # |
| 108 | + # You may find torch.nn.functional's binary_cross_entropy function useful. |
| 109 | + # |
| 110 | + # Return a single value accumulating the loss over the whole batch. |
| 111 | + # |
| 112 | + # Arguments: |
| 113 | + # x is the batch of observations |
| 114 | + # recon_x, mu, and logvar are the outputs of forward(x) (above) - see the usage below |
| 115 | + x = x.view(-1,784) # Reshape x to provide suitable inputs to the encoder |
| 116 | + pass |
| 117 | + |
| 118 | +def train(epoch): |
| 119 | + model.train() |
| 120 | + train_loss = 0 |
| 121 | + for batch_idx, (data, _) in enumerate(train_loader): |
| 122 | + data = data.to(device) |
| 123 | + optimizer.zero_grad() |
| 124 | + recon_batch, mu, logvar = model(data) |
| 125 | + loss = loss_function(recon_batch, data, mu, logvar) |
| 126 | + loss.backward() |
| 127 | + train_loss += loss.item() |
| 128 | + optimizer.step() |
| 129 | + if batch_idx % args.log_interval == 0: |
| 130 | + print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format( |
| 131 | + epoch, batch_idx * len(data), len(train_loader.dataset), |
| 132 | + 100. * batch_idx / len(train_loader), |
| 133 | + loss.item() / len(data))) |
| 134 | + |
| 135 | + print('====> Epoch: {} Average loss: {:.4f}'.format( |
| 136 | + epoch, train_loss / len(train_loader.dataset))) |
| 137 | + |
| 138 | + |
| 139 | +def test(epoch): |
| 140 | + model.eval() |
| 141 | + test_loss = 0 |
| 142 | + with torch.no_grad(): |
| 143 | + for i, (data, _) in enumerate(test_loader): |
| 144 | + data = data.to(device) |
| 145 | + recon_batch, mu, logvar = model(data) |
| 146 | + test_loss += loss_function(recon_batch, data, mu, logvar).item() |
| 147 | + if i == 0: |
| 148 | + n = min(data.size(0), 8) |
| 149 | + comparison = torch.cat([data[:n], |
| 150 | + recon_batch.view(args.batch_size, 1, 28, 28)[:n]]) |
| 151 | + save_image(comparison.cpu(), |
| 152 | + 'results/reconstruction_' + str(epoch) + '.png', nrow=n) |
| 153 | + |
| 154 | + test_loss /= len(test_loader.dataset) |
| 155 | + print('====> Test set loss: {:.4f}'.format(test_loss)) |
| 156 | + |
| 157 | +if __name__ == "__main__": |
| 158 | + for epoch in range(1, args.epochs + 1): |
| 159 | + train(epoch) |
| 160 | + test(epoch) |
| 161 | + with torch.no_grad(): |
| 162 | + sample = torch.randn(64, 20).to(device) |
| 163 | + sample = model.decode(sample).cpu() |
| 164 | + save_image(sample.view(64, 1, 28, 28), |
| 165 | + 'results/sample_' + str(epoch) + '.png') |
0 commit comments