<jupyter_start><jupyter_text><jupyter_code>! nvidia-smi
! git clone https://github.com/satyajitghana/TSAI-DeepVision-EVA4.0-Phase-2<jupyter_output><empty_output><jupyter_text>MNIST VAE<jupyter_code>import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import torchvision.transforms as T
from torchvision.datasets import MNIST
import torchvision.utils as vutils

import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

from pathlib import Path
from PIL import Image

from torchsummary import summary
sns.set()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")<jupyter_output><empty_output><jupyter_text>Build the Models<jupyter_code>class Encoder(nn.Module):
    """
    Creates the Encoder

    nc (int): number of channels of the input image
    ndf (int): number of feature maps for the discriminator

    This uses the special Spectral Normalization ref: https://arxiv.org/abs/1802.05957
    """
    def __init__(self, nc: int = 1, ndf: int = 16, z_dim: int = 20):
        super(Encoder, self).__init__()

        self.main = nn.Sequential(
            # input is (nc) x 28 x 28
            nn.Conv2d(nc, ndf * 2, 3, 1, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.1, inplace=True),

            nn.Conv2d(ndf * 2, ndf * 2, 3, 1, 1, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.1, inplace=True),

            nn.Conv2d(ndf * 2, ndf * 2, 4, 2, 2, bias=False),
            nn.BatchNorm2d(ndf * 2),
            nn.LeakyReLU(0.1, inplace=True),

            # state size. (ndf*4) x 14 x 14
            nn.Conv2d(ndf * 2, ndf * 4, 3, 1, 1, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.1, inplace=True),

            nn.Conv2d(ndf * 4, ndf * 4, 4, 2, 2, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.1, inplace=True),

            # state size. (ndf*4) x 7 x 7
            nn.Conv2d(ndf * 4, ndf * 4, 3, 1, 0, bias=False),
            nn.BatchNorm2d(ndf * 4),
            nn.LeakyReLU(0.1, inplace=True),
            # state size. (ndf*4) x 5 x 5
        )

        self.gap = nn.AdaptiveAvgPool2d(1)

        self.mu = nn.Conv2d(ndf * 4, z_dim, 1, bias=False)
        self.logvar = nn.Conv2d(ndf * 4, z_dim, 1, bias=False)

    def forward(self, input):
        out = self.main(input)
        out = self.gap(out)
        mu = self.mu(out)
        logvar = self.logvar(out)
        return mu, logvar
encoder = Encoder(nc=1, ndf=8).to(device)
input = torch.randn(1, 1, 28, 28).to(device)
encoder(input)[0].shape, encoder(input)[1].shape
summary(encoder, (1, 28, 28))
class Decoder(nn.Module):
    """
    Creates the Decoder

    nc (int): channels in the final output
    nz (int): size of the latent z vector
    ngf (int): number of feature maps for the generator
    """
    def __init__(self, nc: int = 1, ngf: int = 16, z_dim: int = 20):
        super(Decoder, self).__init__()

        self.main = nn.Sequential(
            # input is Z, going into a convolution
            nn.ConvTranspose2d( z_dim, ngf * 8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(ngf * 8),
            nn.ReLU(True),
            # state size. (ngf*4) x 4 x 4
            nn.ConvTranspose2d( ngf * 8, ngf * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            nn.Conv2d(ngf * 4, ngf * 4, 3, 1, 1, bias=False),
            nn.BatchNorm2d(ngf * 4),
            nn.ReLU(True),
            # state size. (ngf*3) x 8 x 8
            nn.ConvTranspose2d( ngf * 4, ngf * 2, 4, 2, 2, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            nn.Conv2d(ngf * 2, ngf * 2, 3, 1, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            # state size. (ngf*2) x 14 x 14
            nn.ConvTranspose2d( ngf * 2, ngf * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(ngf * 2),
            nn.ReLU(True),
            nn.Conv2d(ngf * 2, nc, 3, 1, 1, bias=False),
            # state size. (ngf*2) x 28 x 28
            nn.Sigmoid()
        )

    def forward(self, input):
        return self.main(input)
decoder = Decoder(nc=1, ngf=8, z_dim=20).to(device)
input = torch.randn(1, 20, 1, 1).to(device)
decoder(input).shape
summary(decoder, (20, 1, 1))
class MNISTVAE(nn.Module):
    def __init__(self, nc: int = 1, nf: int = 8, z_dims: int = 20):
        super(MNISTVAE, self).__init__()

        self.encoder = Encoder(nc, nf, z_dims)
        self.decoder = Decoder(nc, nf, z_dims)

    def reparameterize(self, mu, logvar):
        if self.training:
            std = torch.exp(0.5 * logvar)
            eps = torch.randn_like(std)
            return eps.mul(std).add_(mu)
        else:
            return mu
    
    def reparam_like_train(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return eps.mul(std).add_(mu)

    def forward(self, input):
        mu, logvar = self.encoder(input)
        z = self.reparameterize(mu, logvar)
        out = self.decoder(z)

        return out, mu, logvar

mnistvae = MNISTVAE(nc=1, nf=24, z_dims=20).to(device)
input = torch.randn(1, 1, 28, 28).to(device)
[x.shape for x in mnistvae(input)]
summary(mnistvae, (1, 28, 28))
def kld(mu, logvar):
    KLD = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
    return KLD

def loss_function(y_pred, y_true):

    recon_x, mu, logvar = y_pred

    BCE = F.binary_cross_entropy(recon_x, y_true, reduction='sum')

    B, C, H, W = recon_x.size()
    KLD = kld(mu, logvar)
    KLD /= B

    return BCE + KLD
input = torch.randn(32, 1, 28, 28).to(device)
y_true = torch.randn(32, 1, 28, 28).to(device)
output = mnistvae(input)
loss = loss_function(output, y_true)
loss<jupyter_output><empty_output><jupyter_text>Dataset<jupyter_code>train_loader = torch.utils.data.DataLoader(
    MNIST(root='data', train=True, transform=T.ToTensor(), download=True),
    batch_size=128,
    shuffle=True,
    num_workers=2,
    pin_memory=True,
)
# Plot some training images
plt.ioff()
real_batch, _ = next(iter(train_loader))
fig = plt.figure(figsize=(10, 10))
plt.title("Training Images")
for idx, image in enumerate(real_batch[:64]):
    plt.subplot(8, 8, idx+1)
    image.squeeze_(0)
    plt.imshow(image.cpu(), cmap='gray')
    plt.axis('off')
    plt.grid('off')
# plt.show()
fig.suptitle('Train Images')
plt.show()
# fig.tight_layout()
# fig.savefig('test.jpg')
# plt.close()
def save_images(data, title='', dest='test.jpg'):
    plt.ioff()
    fig = plt.figure(figsize=(10, 10))
    # fetch the first 64 images only
    for idx, image in enumerate(data[:64]):
        plt.subplot(8, 8, idx+1)
        image.squeeze_(0)
        plt.imshow(image.cpu(), cmap='gray')
        plt.axis('off')
        plt.grid('off')
    fig.suptitle(title)
    fig.savefig(dest)
    plt.close()<jupyter_output><empty_output><jupyter_text>Train the Model<jupyter_code>device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = MNISTVAE(nc=1, nf=32, z_dims=10)
optimizer = optim.Adam(model.parameters(), lr=1e-3)
from tqdm.auto import tqdm

def train_epoch(epoch, model, optimizer):
    model.train()

    model = model.to(device)

    train_loss = 0
    for batch_idx, (data, _) in enumerate(tqdm(train_loader)):
        data = data.to(device)

        optimizer.zero_grad() # zero out the gradients

        output = model(data) # forward pass

        loss = loss_function(output, data) # compute the loss

        loss.backward() # backprop the loss

        train_loss += loss.item()

        optimizer.step() # update the optimizer
    
    print(f'=> Epoch: {epoch}\tAverage Loss: {(train_loss / len(train_loader)):.4f}\n')
NUM_EPOCHS = 10
out_images = Path('vae_output')
out_images.mkdir(parents=True, exist_ok=True)
for epoch in range(1, NUM_EPOCHS + 1):
    train_epoch(epoch, model, optimizer)

    # with torch.no_grad():
    #     model.eval()
    #     # `64 x z_dims` random array
    #     sample = torch.randn(64, 20, 1, 1).to(device)
    #     sample_out = model.decoder(sample)
    #     sample_out = torch.sigmoid(sample_out)

    # save_images(sample_out, f'Epoch {epoch:02}', dest=f'vae_output/epoch_{epoch:02}')<jupyter_output><empty_output><jupyter_text>Let's try to see what the model outputs if we provide it with some random z vector<jupyter_code>with torch.no_grad():
    model.eval()
    sample = torch.randn(64, 10, 1, 1).to(device)
    output = model.decoder(sample)
fig = plt.figure(figsize=(10, 10))
for idx, image in enumerate(output[:64]):
    plt.subplot(8, 8, idx+1)
    image.squeeze_(0)
    plt.imshow(image.cpu())
    plt.axis('off')
    plt.grid('off')

plt.show()<jupyter_output><empty_output><jupyter_text>Lets see how the model reconstructs input images<jupyter_code>batch, _ = next(iter(train_loader))
with torch.no_grad():
    batch = batch.to(device)

    out, mu, logvar = model(batch)
fig = plt.figure(figsize=(10, 10))
for idx, image in enumerate(out[:64]):
    plt.subplot(8, 8, idx+1)
    image.squeeze_(0)
    plt.imshow(image.cpu())
    plt.axis('off')
    plt.grid('off')
plt.tight_layout()
fig.canvas.draw()
predicted = np.array(fig.canvas.renderer.buffer_rgba())
fig = plt.figure(figsize=(10, 10))
for idx, image in enumerate(batch[:64]):
    plt.subplot(8, 8, idx+1)
    image.squeeze_(0)
    plt.imshow(image.cpu())
    plt.axis('off')
    plt.grid('off')
plt.tight_layout()
fig.canvas.draw()
real = np.array(fig.canvas.renderer.buffer_rgba())
fig = plt.figure(figsize=(15, 10))

plt.subplot(1, 2, 1)
plt.imshow(real)
plt.title('Real')
plt.grid('off')
plt.axis('off')

plt.subplot(1, 2, 2)
plt.imshow(predicted)
plt.title('Predicted')
plt.grid('off')
plt.axis('off')

plt.show()<jupyter_output><empty_output><jupyter_text>Playing with the z vector of the reconstructed images<jupyter_code>batch[0].shape
with torch.no_grad():
    mu, logvar = model.encoder(batch[3].unsqueeze_(0).to(device))
    z = model.reparameterize(mu, logvar)
    z[0][1][0] += 1.5
    z[0][0][0] += -1.5
    out = model.decoder(z)
    out = out.squeeze_(0)
plt.imshow(out[0].cpu())
plt.axis('off')
plt.grid('off')
plt.show()
torch.save(model, 'mnist_vae.pt')
torch.onnx.export(model.encoder,          
                  torch.randn(1, 1, 28, 28).to(device),              
                  "mnist_vae_encoder.onnx",  
                  export_params=True,        
                  opset_version=12,        
                  do_constant_folding=True,  
                  input_names = ['input'],   
                  output_names = ['output'], 
                  dynamic_axes={'input' : {0 : 'batch_size'},  
                                'output' : {0 : 'batch_size'}})
torch.onnx.export(model.decoder,          
                  torch.randn(1, 10, 1, 1).to(device),              
                  "mnist_vae_decoder.onnx",  
                  export_params=True,        
                  opset_version=12,        
                  do_constant_folding=True,  
                  input_names = ['input'],   
                  output_names = ['output'], 
                  dynamic_axes={'input' : {0 : 'batch_size'},  
                                'output' : {0 : 'batch_size'}})<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14663354918832608, 'fraction_numerical': 0.03308790916348798, 'mean_word_length': 2.5677215189873417, 'pattern_counts': {'":': 0, '<': 33, '<?xml version=': 0, '>': 34, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 1, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '8', 'dir': 'jupyter-structured-clean-dedup', 'id': '404614', 'license': 'MIT', 'n_tokens_mistral': 4493, 'n_tokens_neox': 4141, 'n_words': 912, 'path': '07-VariationalAutoEncoders/RedCar_VAE.ipynb', 'repo_name': 'amitkml/TSAI-DeepVision-EVA4.0-Phase-2', 'size': '11273'}