Add pytorch version

Author: rusaha1

pytorch/vae/model.py

@@ -0,0 +1,96 @@
+import torch 
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.autograd import Variable
+
+class VaeNet(nn.Module):
+    """
+    This class will have the encoder and the decoder networks
+    of the variational autoencoder. The encoder will take the 
+    input image and and transform the image to it's latent space.
+    There will also be a separate method related to the ELBO loss
+    which will be calculated as follows:
+
+    Elbo(x) = Marginal_likelihood(x) - KL_divergence(posterior || true_prior)
+
+    Applying the same model to both the MNIST dataset and the
+    CIFAR10 dataset. So the things should be agnostic with each 
+    others those that deal with the dimensionality of things.
+
+    Since we are using the same model for both the datasets, we
+    should keep the capasity as high as possible so that it is
+    capable enough to deal with the complexity of the tougher
+    dataset i.e. CIFAR10.
+    """
+
+    def __init__(self, latent_dim, batch_size):
+        super(VaeNet, self).__init__() 
+        
+        self.latent_dim = latent_dim
+        self.batch_size = batch_size
+        self.encoder_network()
+        self.decoder_network()
+        
+
+    def encoder_network(self):
+        # The definition of the convolutional layer takes in the
+        # number of channels of the input matrix, the number of 
+        # channels for the output matrix, the size of the 
+        # convolutional kernel.
+        self.en_conv1 = nn.Conv2d(3, 8, 3, padding=1) 
+        self.en_conv2 = nn.Conv2d(8, 16, 3)
+        self.en_bn2 = nn.BatchNorm2d(16)
+        self.en_conv3 = nn.Conv2d(16, 32, 2, stride=2) # [BSx32x15x15]
+        self.en_bn3 = nn.BatchNorm2d(32)
+        self.en_conv4 = nn.Conv2d(32, 64, 3, stride=2) 
+        self.en_bn4 = nn.BatchNorm2d(64)
+        self.en_conv5 = nn.Conv2d(64, 128, 3, stride=2) # [BSx128x3x3]
+        self.en_bn5 = nn.BatchNorm2d(128)
+        self.en_fc1 = nn.Linear(128 * 3 * 3, 256)
+        self.en_mu = nn.Linear(256, self.latent_dim)
+        self.en_sigma = nn.Linear(256, self.latent_dim)
+
+    def decoder_network(self):
+        self.de_deconv1 = nn.ConvTranspose2d(self.latent_dim, self.batch_size * 4, 4, 1, 0)
+        self.de_bn1 = nn.BatchNorm2d(self.batch_size * 4)
+        self.de_deconv2 = nn.ConvTranspose2d(self.batch_size * 4, self.batch_size * 2, 4, 2, 1)
+        self.de_bn2 = nn.BatchNorm2d(self.batch_size * 2)
+        self.de_deconv3 = nn.ConvTranspose2d(self.batch_size * 2, self.batch_size, 4, 2, 1)
+        self.de_bn3 = nn.BatchNorm2d(self.batch_size)
+        self.de_deconv4 = nn.ConvTranspose2d(self.batch_size, 3, 4, 2, 1)
+
+    def encoder(self, x):
+        x = F.elu(self.en_conv1(x))
+        x = F.elu(self.en_bn2(self.en_conv2(x)))
+        x = F.elu(self.en_bn3(self.en_conv3(x)))
+        x = F.elu(self.en_bn4(self.en_conv4(x)))
+        x = F.elu(self.en_bn5(self.en_conv5(x)))
+        x = x.view(-1, 128 * 3 *3) # flatten
+        x = F.elu(self.en_fc1(x))
+        x_mu = self.en_mu(x)
+        x_log_sigma_sq = self.en_sigma(x)
+        return x_mu, x_log_sigma_sq
+
+    def reparameterize(self, mu, log_sigma_sq):
+        if self.training:
+            std = log_sigma_sq.mul(0.5).exp_() # Doing things in place
+            eps = Variable(std.data.new(std.size()).normal_())
+            return eps.mul(std).add_(mu) # multiply the std with epsilon and add it to the mean
+        else:
+            return mu
+
+    def decoder(self, x):
+        x = x.view(-1, self.latent_dim, 1, 1)
+        x = F.elu(self.de_bn1(self.de_deconv1(x)))
+        x = F.elu(self.de_bn2(self.de_deconv2(x)))
+        x = F.elu(self.de_bn3(self.de_deconv3(x)))
+        x = F.tanh((self.de_deconv4(x)))
+        return x
+    
+    def forward(self,x):
+       mu, log_sigma_sq = self.encoder(x)
+       z = self.reparameterize(mu, log_sigma_sq)
+       reconstructed_img = self.decoder(z)
+       return reconstructed_img, mu, log_sigma_sq
+
+

pytorch/vae/trainer.py

@@ -0,0 +1,125 @@
+from __future__ import print_function
+import torch
+import torch.utils.data
+from torch import nn, optim
+from torch.autograd import Variable
+import argparse
+import numpy as np
+from torch.nn import functional as F
+from torchvision import datasets, transforms
+from torchvision.utils import save_image
+from model import VaeNet
+
+# Define arguments required for training using parser.
+parser = argparse.ArgumentParser(description='VAE CIFAR 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',
+                    help='number of epochs to train (default: 10)')
+parser.add_argument('--no_cuda', action='store_true', default=False,
+                    help='enables CUDA training')
+parser.add_argument('--seed', type=int, default=1, metavar='S',
+                    help='random seed (default: 1)')
+parser.add_argument('--latent_dim', type=int, default=100, metavar='L',
+                    help='size of the latent dimension (default: 100)')
+parser.add_argument('--log_interval', type=int, default=100, metavar='N',
+                    help='how many batches to wait for before logging training status')
+
+# Parse the arguments and see if cuda is available
+args = parser.parse_args()
+args.cuda = not args.no_cuda and torch.cuda.is_available()
+
+# Use the defined seed to initialize state
+torch.manual_seed(args.seed)
+if args.cuda:
+    torch.cuda.manual_seed(args.seed)
+
+# Define the transformation process of the data.
+transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
+
+trainset = datasets.CIFAR10(root='./data/', train=True, download=True,
+        transform=transform)
+trainloader = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size, shuffle=True,
+                                            num_workers=3)
+testset = datasets.CIFAR10(root='./data/', train=False, download=True, 
+                                        transform=transform)
+testloader = torch.utils.data.DataLoader(testset, batch_size=args.batch_size, shuffle=False,
+                                            num_workers=3)
+# Define the model and port it to the gpu
+model = VaeNet(batch_size=args.batch_size, latent_dim=args.latent_dim)
+if args.cuda:
+    model = model.cuda()
+
+optimizer = optim.Adam(model.parameters(), lr=1e-4)
+
+# Define the loss function.
+def vae_loss(x_recons, x_original, mu, log_sigma_sq):
+    reconstruct_loss = F.mse_loss(x_recons, x_original)
+    # KL divergence loss can be defined as follows
+    # 0.5 * sum(1 + log(sigma^2) - mu^2 -sigma^2)
+    kl_div = -0.5 * torch.sum(1 + log_sigma_sq - mu.pow(2) - log_sigma_sq.exp())
+    kl_div /= args.batch_size * 32 * 32 * 3
+    return kl_div, reconstruct_loss
+
+# Define the train step
+def train(epoch):
+    model.train()
+    train_loss = 0
+    likelihood = 0
+    divergence = 0
+    for batch_idx, data in enumerate(trainloader):
+        images, labels = data
+        if args.cuda:
+            images = Variable(images.cuda())
+        optimizer.zero_grad()
+        reconstructed_img, mu, log_sigma_sq = model(images)
+        kl_div, recon_loss = vae_loss(reconstructed_img, images, mu, log_sigma_sq)
+        loss = kl_div + recon_loss
+        loss.backward()
+        train_loss += loss.data[0]
+        likelihood += recon_loss
+        divergence += kl_div
+        optimizer.step()
+        if batch_idx % args.log_interval == 0:
+            print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
+                epoch, batch_idx * len(data), len(trainloader.dataset),
+                100. * batch_idx / len(trainloader),
+                loss.data[0] / len(data)))
+    print('Epoch: %f, Average loss: %f, Average reconstruction loss: %f, Average kl divergence loss: %f' \
+           % (epoch, train_loss / len(trainloader.dataset), \
+           likelihood / len(trainloader.dataset), \
+           divergence / len(trainloader.dataset)))
+
+# Define the test step
+def test(epoch):
+    model.eval()
+    test_loss = 0
+    for batch_idx, data in enumerate(testloader):
+        images, labels = data
+        if args.cuda:
+            images = Variable(images.cuda())
+        reconstructed_img, mu, log_sigma_sq = model(images)
+        kl_div, recon_loss = vae_loss(reconstructed_img, images, mu, log_sigma_sq)
+        test_loss += (kl_div + recon_loss).data[0]
+        if batch_idx == 0:
+            n = min(images.size(0), 8)
+            comparison = torch.cat([images[:n], 
+                                   reconstructed_img.view(args.batch_size, 3, 32, 32)[:n]])
+            save_image(comparison.data, 
+                        'results/reconstruction_' + str(epoch) + '.png', nrow=n)
+
+    test_loss /= len(testloader.dataset)
+    print('Test set loss: %f' % (test_loss))
+
+# Set up the training loop
+for epoch in range(1, args.epochs + 1):
+    train(epoch)
+    test(epoch)
+    # Sample a random value from the gaussian distribution
+    sample = Variable(torch.randn(args.batch_size, 100))
+    if args.cuda:
+        sample = sample.cuda()
+    sample = model.decoder(sample)
+    save_image(sample.data.view(args.batch_size, 3, 32, 32), 
+            'results/sample_' + str(epoch) + '.png')
+