vae_mnist_cnn.py

# -*- coding: utf-8 -*-
"""vae_mnist_cnn.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1FivAIGGFfx8J5R5-2aRbx5PWiWxIwI3j

### Loading Required Libraries
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
import torch.optim as optim

from torch.utils.data import DataLoader
from torchvision import datasets, transforms

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from sklearn.manifold import TSNE

from sklearn.metrics import accuracy_score
from sklearn.metrics import f1_score
from sklearn.metrics import confusion_matrix
from sklearn.metrics import recall_score
from sklearn.metrics import precision_score

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

BATCH_SIZE = 64         # number of data points in each batch
N_EPOCHS = 20           # times to run the model on complete data
INPUT_DIM = 28 * 28     # size of each input
HIDDEN_DIM = 256        # hidden dimension
LATENT_DIM = 10         # latent vector dimension
lr = 1e-3               # learning rate

"""### Loading Data"""

transforms = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(
    './data',
    train=True,
    download=True,
    transform=transforms)

test_dataset = datasets.MNIST(
    './data',
    train=False,
    download=True,
    transform=transforms
)

train_iterator = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)
test_iterator = DataLoader(test_dataset, batch_size=BATCH_SIZE)

"""### Batch Visualization of Given Data"""

def visualize_data(batch):
    batch = torchvision.utils.make_grid(batch)
    batch = batch.numpy()
    batch = np.transpose(batch,(1,2,0))
    plt.figure(figsize = (8,8))
    plt.imshow(batch, cmap = 'Greys_r')
    plt.show()

batch,labels = iter(train_iterator).next()
visualize_data(batch)

"""### Creating Model"""

class Encoder(nn.Module):
    ''' This the encoder part of VAE
    '''
    def __init__(self, hidden_dim, z_dim):
        '''
        Args:
            input_dim: A integer indicating the size of input (in case of MNIST 28 * 28).
            hidden_dim: A integer indicating the size of hidden dimension.
            z_dim: A integer indicating the latent dimension.
        '''
        super().__init__()
        self.conv1 = nn.Conv2d(in_channels=1, out_channels = 16, kernel_size=3)
        self.conv2 = nn.Conv2d(in_channels = 16, out_channels = 32, kernel_size = 3)

        self.linear = nn.Linear(32*24*24, 512)
        self.linear2 = nn.Linear(512, hidden_dim)

        self.mu = nn.Linear(hidden_dim, z_dim)
        self.var = nn.Linear(hidden_dim, z_dim)

    def forward(self, x):
        # x is of shape [batch_size, input_dim]
        x = F.relu(self.conv1(x))
        x = F.relu(self.conv2(x))
        bs, ch, w,h = x.shape
        x = x.view(bs,ch*w*h)


        x = F.relu(self.linear(x))
        hidden = F.relu(self.linear2(x))

        # hidden is of shape [batch_size, hidden_dim]
        z_mu = self.mu(hidden)
        # z_mu is of shape [batch_size, latent_dim]
        z_var = self.var(hidden)
        # z_var is of shape [batch_size, latent_dim]

        return z_mu, z_var

class Decoder(nn.Module):
    ''' This the decoder part of VAE
    '''
    def __init__(self, z_dim, hidden_dim):
        
        super().__init__()

        self.linear = nn.Linear(z_dim, hidden_dim)
        self.linear2 = nn.Linear(hidden_dim, 512 )
        self.linear3 = nn.Linear(512, 18432)

        self.deconv1 = nn.ConvTranspose2d(32, 16, kernel_size=3)
        self.deconv2 = nn.ConvTranspose2d(16, 1, kernel_size = 3)

    def forward(self, x):
        # x is of shape [batch_size, latent_dim]
        x = F.relu(self.linear(x))
        x = F.relu(self.linear2(x))
        x = F.relu(self.linear3(x))
        # print(x.shape)
        x = x.view(-1, 32, 24,24)

        x = F.relu(self.deconv1(x))
        predicted = torch.sigmoid(self.deconv2(x))
        # predicted is of shape [batch_size, output_dim]
        return predicted

class VAE(nn.Module):
    def __init__(self, enc, dec):
        ''' This the VAE, which takes a encoder and decoder.
        '''
        super().__init__()

        self.enc = enc
        self.dec = dec

    def forward(self, x):
        # encode
        z_mu, z_var = self.enc(x)

        # sample from the distribution having latent parameters z_mu, z_var
        # reparameterize
        std = torch.exp(z_var / 2)
        eps = torch.randn_like(std)
        x_sample = eps.mul(std).add_(z_mu)

        # decode
        predicted = self.dec(x_sample)
        return predicted, z_mu, z_var

# encoder
encoder = Encoder( HIDDEN_DIM, LATENT_DIM)

# decoder
decoder = Decoder(LATENT_DIM, HIDDEN_DIM)

# vae
model = VAE(encoder, decoder).cuda()

# optimizer
optimizer = optim.Adam(model.parameters(), lr=lr)

# xavier weight initialization
def init_xavier(m):
    if type(m) == nn.Linear:
        nn.init.xavier_uniform(m.weight)

model.apply(init_xavier)

print(model)

"""### Training the Model"""

def train():
    # set the train mode
    model.train()

    # loss of the epoch
    train_loss = 0

    for i, (x, _) in enumerate(train_iterator):

        x = x.to(device)

        # update the gradients to zero
        optimizer.zero_grad()

        # forward pass
        x_sample, z_mu, z_var = model(x)

        # reconstruction loss
        recon_loss = F.binary_cross_entropy(x_sample, x, size_average=False)

        # kl divergence loss
        kl_loss = 0.5 * torch.sum(torch.exp(z_var) + z_mu**2 - 1.0 - z_var)

        # total loss
        loss = recon_loss + kl_loss

        # backward pass
        loss.backward()
        train_loss += loss.item()

        # update the weights
        optimizer.step()

    return train_loss

def test():
    # set the evaluation mode
    model.eval()

    # test loss for the data
    test_loss = 0

    with torch.no_grad():
        for i, (x, _) in enumerate(test_iterator):

            x = x.to(device)

            # forward pass
            x_sample, z_mu, z_var = model(x)

            # reconstruction loss
            recon_loss = F.binary_cross_entropy(x_sample, x, size_average=False)

            # kl divergence loss
            kl_loss = 0.5 * torch.sum(torch.exp(z_var) + z_mu**2 - 1.0 - z_var)

            # total loss
            loss = recon_loss + kl_loss
            test_loss += loss.item()

    return test_loss

def generate_data(bs=BATCH_SIZE):
    z = torch.randn(bs, LATENT_DIM).to(device)
    reconstructed_img = model.dec(z)
    img = reconstructed_img.view(-1,1,28, 28).data

    img = img.cpu()
    visualize_data(img)

N_EPOCHS = 20

def fit():
    best_test_loss = float('inf')
    print('Before Training : ')
    generate_data()

    for e in range(N_EPOCHS):

        train_loss = train()
        test_loss = test()

        train_loss /= len(train_dataset)
        test_loss /= len(test_dataset)

        print(f'Epoch {e}, Train Loss: {train_loss:.2f}, Test Loss: {test_loss:.2f}')
        generate_data()

        if best_test_loss > test_loss:
            best_test_loss = test_loss
            patience_counter = 1
        else:
            patience_counter += 1

        if patience_counter > 3:
            break

"""### Saving The Weights of Best Model"""

# fit()      Training The model

# path = '/content/drive/My Drive/Colab Notebooks/deeplearning/asg3/pickle/best_vae_model_state_dict.pt'
# torch.save(model.state_dict(), path)

# Loading the weights of model
# model = VAE(encoder, decoder)

path = '/content/drive/My Drive/Colab Notebooks/deeplearning/asg3/pickle/best_vae_model_state_dict.pt'
model.load_state_dict(torch.load(path))

"""### Generating grid of data Using VAE Model"""

generate_data() # This function defined above will generate Batch Data

"""### Get Latent data"""

def get_latent_data(data_iterator):

    means = np.zeros((1,10))
    vars = np.zeros((1,10))

    data_labels = []
    with torch.no_grad():
        for batch, labels in data_iterator:
            labels = list(labels.numpy())
            batch = batch.to(device)

            mean, var = model.enc(batch)

            mean = mean.to('cpu').numpy()
            var = var.to('cpu').numpy()

            means = np.vstack((means,mean))
            vars = np.vstack((vars, var))

            data_labels.extend(labels)

    means = means[1:]
    vars = vars[1:]

    return means, vars, data_labels

train_means, train_vars, train_labels = get_latent_data(train_iterator)

"""### TSNE"""

mean_tsne = TSNE(n_components = 1, random_state = 0)
var_tsne = TSNE(n_components = 1, random_state = 0)

tsne_mean = mean_tsne.fit_transform(train_means)
tsne_var = var_tsne.fit_transform(train_vars)


df = pd.DataFrame()
df['Mean'] = tsne_mean.squeeze(1)
df['Variance'] = tsne_var.squeeze(1)
df['label'] = train_labels

g = sns.FacetGrid(df, hue = 'label', height = 8).map(plt.scatter, 'Mean','Variance').add_legend()
plt.legend(fontsize='x-large', title_fontsize='40',loc = 'best')

plt.show()

"""### Getting latent data for test data"""

train_latent_data = np.column_stack((train_means, train_vars))
train_latent_df = pd.DataFrame(train_latent_data)  
train_labels = np.array(train_labels)

# test latent data 
test_means, test_vars, test_labels = get_latent_data(test_iterator)

test_latent_data = np.column_stack((test_means, test_vars))
test_latent_df = pd.DataFrame(test_latent_data)  
test_labels = np.array(test_labels)

print('Shape of Training data : ', train_latent_data.shape)
print('Shape of Testing Data : ', test_latent_data.shape)

"""### Training the SVM classifier on Latent Space"""

from sklearn.svm import LinearSVC

clf = LinearSVC(random_state=0, tol=1e-5)
clf.fit(train_latent_df, train_labels)

# prediction on test data

ypreds = clf.predict(test_latent_df)


def plot_confusion_matrix(cm, title, filepath):

    df_cm = pd.DataFrame(cm, index = [i for i in "0123456789"],
                      columns = [i for i in "0123456789"])

    plt.figure(figsize = (8,5))
    sns.heatmap(df_cm, annot=True, fmt='d')
    plt.title(title)
    #plt.savefig(filepath)
    plt.show()

test_acc = accuracy_score(ytrue, ypreds)
f_score = f1_score(ytrue, ypreds, average='macro')

prec = precision_score(ytrue, ypreds, average='macro')
recall = recall_score(ytrue, ypreds, average='macro')

cm = confusion_matrix(ytrue, ypreds)

print('Test Accuracy : ', test_acc)
print('Test prec : ', prec)
print('Test Recall : ', recall)
print('Test Data F - Score : ', f_score)

plot_confusion_matrix(cm, 'Confusion Matrix ','dummy')