<jupyter_start><jupyter_text>Deep convolutional generative adversarial network (GAN) for generating MNIST-like images of handwritten digits Before we start we need to load libraries and setup the model hyperparameters<jupyter_code>import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torchvision import datasets, transforms

import numpy as np

import matplotlib.pyplot as plt
plt.rcParams.update({'font.size': 14})

model_args = {}
# random seed
model_args['seed'] = 123
# we will use batch size of 128 in Stochastic Gradient Descent (SGD) optimization of the network
model_args['batch_size'] = 128 
# learning rate is how fast it will descend
model_args['lr'] = 2.e-4
# the number of epochs is the number of times you go through the full dataset 
model_args['epochs'] = 50
# GAN parameters:
model_args['latent_space_dim'] = 100<jupyter_output><empty_output><jupyter_text>Let's start with loading the dataset and definig dataloaders<jupyter_code># load the MINST dataset via torchvision 
transform = transforms.ToTensor()
train_subset = datasets.MNIST('../data', train=True, download=True, transform=transform)
# define dataloaders
loader_kwargs = {'batch_size': model_args['batch_size'], 
                 'num_workers': 2, 
                 'pin_memory': True, 
                 'shuffle': True}
train_loader = torch.utils.data.DataLoader(train_subset, **loader_kwargs)<jupyter_output><empty_output><jupyter_text>Plot sample images<jupyter_code>example_number = 123
fig, axs = plt.subplots(5, 5, figsize=(7,7), tight_layout=True)
for i in range(5):
    for j in range(5):
        axs[i,j].imshow(train_subset[example_number+i*5+j][0].reshape(28,28), cmap='gray')
        axs[i,j].set_title(train_subset[example_number+i*5+j][1])<jupyter_output><empty_output><jupyter_text>Let's now construct a GAN model:- we are using **nn.Sequential** to define the generator and discriminator models,- generator utilizes *2D transposed convolutions* to size up feature maps in decoder.<jupyter_code>class Flatten(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), -1)
    
class Reshape(nn.Module):
    def forward(self, input):
        return input.view(input.size(0), 64, 7, 7)

class GAN(nn.Module):
    # this defines the generator and discriminator structures
    def __init__(self, latent_space_dim):
        super(GAN, self).__init__()
        self.latent_space_dim = latent_space_dim
        
        self.generator = nn.Sequential(
              
            nn.Linear(self.latent_space_dim, 3136, bias=False),
            nn.BatchNorm1d(num_features=3136),
            nn.LeakyReLU(inplace=True, negative_slope=0.0001),
            Reshape(),
            
            nn.ConvTranspose2d(in_channels=64, out_channels=32, kernel_size=(3, 3), stride=(2, 2), padding=1, bias=False),
            nn.BatchNorm2d(num_features=32),
            nn.LeakyReLU(inplace=True, negative_slope=0.0001),
            #nn.Dropout2d(p=0.2),
            
            nn.ConvTranspose2d(in_channels=32, out_channels=16, kernel_size=(3, 3), stride=(2, 2), padding=1, bias=False),
            nn.BatchNorm2d(num_features=16),
            nn.LeakyReLU(inplace=True, negative_slope=0.0001),
            #nn.Dropout2d(p=0.2),
            
            nn.ConvTranspose2d(in_channels=16, out_channels=8, kernel_size=(3, 3), stride=(1, 1), padding=0, bias=False),
            nn.BatchNorm2d(num_features=8),
            nn.LeakyReLU(inplace=True, negative_slope=0.0001),
            #nn.Dropout2d(p=0.2),
            
            nn.ConvTranspose2d(in_channels=8, out_channels=1, kernel_size=(2, 2), stride=(1, 1), padding=0, bias=False),
            nn.Tanh()
        )
        
        self.discriminator = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=8, padding=1, kernel_size=(3, 3), stride=(2, 2), bias=False),
            nn.BatchNorm2d(num_features=8),
            nn.LeakyReLU(inplace=True, negative_slope=0.0001), 
            #nn.Dropout2d(p=0.2),
            
            nn.Conv2d(in_channels=8, out_channels=32, padding=1, kernel_size=(3, 3), stride=(2, 2), bias=False),
            nn.BatchNorm2d(num_features=32),
            nn.LeakyReLU(inplace=True, negative_slope=0.0001), 
            #nn.Dropout2d(p=0.2),
            
            Flatten(),

            nn.Linear(7*7*32, 1),
            #nn.Sigmoid()
        )

    def generator_forward(self, z):
        return self.generator(z)
    
    def discriminator_forward(self, img):
        return self.discriminator(img).view(-1)<jupyter_output><empty_output><jupyter_text>And define training and plotting utils:<jupyter_code>def train(model, device, train_loader, optimizer_gen, optimizer_dis, epoch_number):
    model.train()
    gen_losses = []
    dis_losses = []
    # get subsequent batches over the data in a given epoch
    for batch_idx, (data, _) in enumerate(train_loader):
        # normalize images to [-1, 1] range
        data = data*2. - 1.
        batch_size = data.size(0)
        # send data tensors to GPU (or CPU)
        data = data.to(device)
        # labels for training discriminator
        valid = torch.ones(batch_size).float().to(device)
        fake = torch.zeros(batch_size).float().to(device)
        # --------------------------
        # Train Generator
        # --------------------------
        # this will zero out the gradients for this batch
        optimizer_gen.zero_grad()
        # generate new images from noise
        z = torch.zeros((batch_size, model_args['latent_space_dim'])).uniform_(-1.0, 1.0).to(device)
        generated_images = model.generator_forward(z)   
        # loss for fooling the discriminator
        dis_predictions = model.discriminator_forward(generated_images.view(batch_size, 1, 28, 28))
        #
        gen_loss = F.binary_cross_entropy_with_logits(dis_predictions, valid) # not fooling D increases G loss
        gen_loss.backward()
        optimizer_gen.step()
        gen_losses.append(gen_loss.item())
        # --------------------------
        # Train Discriminator
        # --------------------------        
        # this will zero out the gradients for this batch
        optimizer_dis.zero_grad()
        # D predictions for real images
        dis_predictions_real = model.discriminator_forward(data.view(batch_size, 1, 28, 28))
        real_loss = F.binary_cross_entropy_with_logits(dis_predictions_real, valid) # D should classify real as valid
        # D predictions for generated images
        dis_predictions_fake = model.discriminator_forward(generated_images.view(batch_size, 1, 28, 28).detach())
        fake_loss = F.binary_cross_entropy_with_logits(dis_predictions_fake, fake) # D should classify generated as fake
        # 
        dis_loss = 0.5*(real_loss + fake_loss)
        dis_loss.backward()
        optimizer_dis.step()        
        dis_losses.append(dis_loss.item())

    return gen_losses, dis_losses


def plot_loss(gen_loss, dis_loss, title=None):
    plt.grid(True)
    plt.xlabel("subsequent epochs")
    plt.ylabel('loss')
    xlabels = np.linspace(0, model_args['epochs'], num=len(gen_loss), endpoint=True)
    plt.plot(xlabels, gen_loss, label='generator')
    plt.plot(xlabels, dis_loss, label='discriminator')
    plt.legend()
    if title is not None: plt.title(title)
    plt.show()<jupyter_output><empty_output><jupyter_text>Experiments Let's define and plot the model<jupyter_code>device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = GAN(model_args['latent_space_dim']).to(device)
print(model)

optimizer_gen = torch.optim.Adam(model.generator.parameters(), lr=model_args['lr'])
optimizer_dis = torch.optim.Adam(model.discriminator.parameters(), lr=model_args['lr'])
from torchsummary import summary
model = model.to(device)
summary(model.generator, input_size=(100,))
summary(model.discriminator, input_size=(1, 28, 28))<jupyter_output>----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Linear-1                 [-1, 3136]         313,600
       BatchNorm1d-2                 [-1, 3136]           6,272
         LeakyReLU-3                 [-1, 3136]               0
           Reshape-4             [-1, 64, 7, 7]               0
   ConvTranspose2d-5           [-1, 32, 13, 13]          18,432
       BatchNorm2d-6           [-1, 32, 13, 13]              64
         LeakyReLU-7           [-1, 32, 13, 13]               0
   ConvTranspose2d-8           [-1, 16, 25, 25]           4,608
       BatchNorm2d-9           [-1, 16, 25, 25]              32
        LeakyReLU-10           [-1, 16, 25, 25]               0
  ConvTranspose2d-11            [-1, 8, 27, 27]           1,152
      BatchNorm2d-12            [-1, 8, 27, 27]              16
        LeakyReLU-13            [-1, 8[...]<jupyter_text>Then train the model<jupyter_code>torch.manual_seed(model_args['seed'])
generator_loss = []
discriminator_loss = []
for epoch_number in range(1, model_args['epochs'] + 1):
    gen_losses, dis_losses = train(model, device, train_loader, optimizer_gen, optimizer_dis, epoch_number)
    print('Train Epoch {} | Generator loss = {:.4f} | Discriminator loss = {:.4f}'.format(
        epoch_number, np.mean(gen_losses), np.mean(dis_losses)))
    generator_loss.extend(gen_losses)
    discriminator_loss.extend(dis_losses)

plot_loss(generator_loss, discriminator_loss)<jupyter_output>Train Epoch 1 | Generator loss = 1.9426 | Discriminator loss = 0.1916
Train Epoch 2 | Generator loss = 1.4060 | Discriminator loss = 0.3513
Train Epoch 3 | Generator loss = 1.3090 | Discriminator loss = 0.3914
Train Epoch 4 | Generator loss = 1.2998 | Discriminator loss = 0.4008
Train Epoch 5 | Generator loss = 1.2033 | Discriminator loss = 0.4513
Train Epoch 6 | Generator loss = 1.1926 | Discriminator loss = 0.4492
Train Epoch 7 | Generator loss = 1.2000 | Discriminator loss = 0.4513
Train Epoch 8 | Generator loss = 1.2080 | Discriminator loss = 0.4598
Train Epoch 9 | Generator loss = 1.1894 | Discriminator loss = 0.4713
Train Epoch 10 | Generator loss = 1.1708 | Discriminator loss = 0.4847
Train Epoch 11 | Generator loss = 1.1477 | Discriminator loss = 0.4909
Train Epoch 12 | Generator loss = 1.1393 | Discriminator loss = 0.4927
Train Epoch 13 | Generator loss = 1.1215 | Discriminator loss = 0.5043
Train Epoch 14 | Generator loss = 1.1077 | Discriminator loss = 0.5143
Train Epoch 15 [...]<jupyter_text>Generate images<jupyter_code>model.eval()

z = torch.zeros((25, model_args['latent_space_dim'])).uniform_(-1.0, 1.0).to(device)
generated_images = model.generator_forward(z)
imgs = generated_images.view(-1, 28, 28)


fig, axs = plt.subplots(5, 5, figsize=(7,7), tight_layout=True)
for i in range(5):
    for j in range(5):
        axs[i,j].imshow(imgs[i*5+j].to(torch.device('cpu')).detach(), cmap='binary')<jupyter_output><empty_output>
StarcoderdataJupyterStructuredCleanDedup