diff --git a/computer_vision/Generative_Adversarial_Network_MNIST.py b/computer_vision/Generative_Adversarial_Network_MNIST.py new file mode 100644 index 000000000000..45a9e6613a35 --- /dev/null +++ b/computer_vision/Generative_Adversarial_Network_MNIST.py @@ -0,0 +1,274 @@ +""" +Generative Adversarial Network + +Objective : To train a GAN model to generate handwritten digits that can be transferred to other domains. + +Resources GAN Theory : + https://en.wikipedia.org/wiki/Generative_adversarial_network +Resources PyTorch: https://pytorch.org/ + +Download dataset from : +PyTorch internal function + +1. Fetch the Dataset with PyTorch function. +2. Create Dataloader. +3. Create Discriminator and Generator. +4. Set the hyperparameters and models. +5. Set the loss functions. +6. Create the training loop. +7. Visualize the losses. +8. Visualize the result from GAN. + +""" + + + + +%matplotlib inline + +import numpy as np +import torch +import matplotlib.pyplot as plt +from torchvision import datasets +import torchvision.transforms as transforms + +# number of subprocesses to use for data loading +num_workers = 0 +# how many samples per batch to load +batch_size = 64 + +# convert data to torch.FloatTensor +transform = transforms.ToTensor() + +# get the training datasets +train_data = datasets.MNIST(root='data', train=True, + download=True, transform=transform) + +# prepare data loader +train_loader = torch.utils.data.DataLoader(train_data, batch_size=batch_size, + num_workers=num_workers) + +import torch.nn as nn +import torch.nn.functional as F + +# Creating Generator and Discriminator for GAN + +class discriminator(nn.Module): + def __init__(self,input_size,output_size,hidden_dim): + super(discriminator,self).__init__() + + #defining the layers of the discriminator + self.fc1 = nn.Linear(input_size,hidden_dim*4) + self.fc2 = nn.Linear(hidden_dim*4,hidden_dim*2) + self.fc3 = nn.Linear(hidden_dim*2,hidden_dim) + #final fully connected layer + self.fc4 = nn.Linear(hidden_dim,output_size) + #dropout layer + self.dropout = nn.Dropout(0.2) + + def forward(self,x): + # pass x through all layers + # apply leaky relu activation to all hidden layers + x = x.view(-1,28*28) #flattening the image + x = F.leaky_relu(self.fc1(x),0.2) + x = self.dropout(x) + x = F.leaky_relu(self.fc2(x),0.2) + x = self.dropout(x) + x = F.leaky_relu(self.fc3(x),0.2) + x = self.dropout(x) + x_out = self.fc4(x) + + return x_out + +class generator(nn.Module): + + def __init__(self, input_size, output_size,hidden_dim): + super(generator, self).__init__() + + # define all layers + self.fc1 = nn.Linear(input_size,hidden_dim) + self.fc2 = nn.Linear(hidden_dim,hidden_dim*2) + self.fc3 = nn.Linear(hidden_dim*2,hidden_dim*4) + #final layer + self.fc4 = nn.Linear(hidden_dim*4,output_size) + #dropout layer + self.dropout = nn.Dropout(0.2) + + + def forward(self, x): + # pass x through all layers + # final layer should have tanh applied + x = F.leaky_relu(self.fc1(x),0.2) + x = self.dropout(x) + x = F.leaky_relu(self.fc2(x),0.2) + x = self.dropout(x) + x = F.leaky_relu(self.fc3(x),0.2) + x = self.dropout(x) + x_out = F.tanh(self.fc4(x)) + return x_out + +# Calculate losses +def real_loss(D_out, smooth=False): + # compare logits to real labels + # smooth labels if smooth=True + #puting it into cuda + batch_size = D_out.size(0) + if smooth: + labels = torch.ones(batch_size).cuda()*0.9 + else: + labels = torch.ones(batch_size).cuda() + criterion = nn.BCEWithLogitsLoss() + loss = criterion(D_out.squeeze(),labels) + return loss + +def fake_loss(D_out): + # compare logits to fake labels + batch_size = D_out.size(0) + labels = torch.zeros(batch_size).cuda() + criterion = nn.BCEWithLogitsLoss() + loss = criterion(D_out.squeeze(),labels) + return loss + +# Discriminator hyperparams +# Size of input image to discriminator (28*28) +input_size = 784 +# Size of discriminator output (real or fake) +d_output_size = 1 +# Size of *last* hidden layer in the discriminator +d_hidden_size = 32 + +# Generator hyperparams + +# Size of latent vector to give to generator +z_size = 100 +# Size of discriminator output (generated image) +g_output_size = 784 +# Size of *first* hidden layer in the generator +g_hidden_size = 32 + +# instantiate discriminator and generator and put it in cuda mode +D = discriminator(input_size, d_output_size,d_hidden_size).cuda() +G = generator(z_size, g_output_size, g_hidden_size).cuda() + +import pickle as pkl + +# training hyperparams +num_epochs = 40 + +# keep track of loss and generated, "fake" samples +samples = [] +losses = [] + +print_every = 400 + +# Get some fixed data for sampling. These are images that are held +# constant throughout training, and allow us to inspect the model's performance +sample_size=16 +fixed_z = np.random.uniform(-1, 1, size=(sample_size, z_size)) +fixed_z = torch.from_numpy(fixed_z).float().cuda() + +# train the network +D.train() +G.train() +for epoch in range(num_epochs): + + for batch_i, (real_images, _) in enumerate(train_loader): + + batch_size = real_images.size(0) + + ## Important rescaling step ## + real_images = (real_images*2 - 1).cuda() # rescale input images from [0,1) to [-1, 1) + + # ============================================ + # TRAIN THE DISCRIMINATOR + # ============================================ + d_optimizer.zero_grad() + # 1. Train with real images + + # Compute the discriminator losses on real images + # use smoothed labels + D_real = D(real_images) + d_real_loss = real_loss(D_real,smooth=True) + # 2. Train with fake images + + # Generate fake images + z = np.random.uniform(-1, 1, size=(batch_size, z_size)) + z = torch.from_numpy(z).float().cuda() + fake_images = G(z) + + # Compute the discriminator losses on fake images + D_fake = D(fake_images) + d_fake_loss = fake_loss(D_fake) + # add up real and fake losses and perform backprop + d_loss = d_real_loss + d_fake_loss + d_loss.backward() + d_optimizer.step() + + + # ========================================= + # TRAIN THE GENERATOR + # ========================================= + g_optimizer.zero_grad() + # 1. Train with fake images and flipped labels + + # Generate fake images + z = np.random.uniform(-1, 1, size=(batch_size, z_size)) + z = torch.from_numpy(z).float().cuda() + fake_images = G(z) + # Compute the discriminator losses on fake images + # using flipped labels! + D_fake = D(fake_images) + # perform backprop + g_loss = real_loss(D_fake) + g_loss.backward() + g_optimizer.step() + + + # Print some loss stats + if batch_i % print_every == 0: + # print discriminator and generator loss + print('Epoch [{:5d}/{:5d}] | d_loss: {:6.4f} | g_loss: {:6.4f}'.format( + epoch+1, num_epochs, d_loss.item(), g_loss.item())) + + + ## AFTER EACH EPOCH## + # append discriminator loss and generator loss + losses.append((d_loss.item(), g_loss.item())) + + # generate and save sample, fake images + G.eval() # eval mode for generating samples + samples_z = G(fixed_z) + samples.append(samples_z) + G.train() # back to train mode + + +# Save training generator samples +with open('train_samples.pkl', 'wb') as f: + pkl.dump(samples, f) + +#ploting Discriminator and Generator loss +fig, ax = plt.subplots() +losses = np.array(losses) +plt.plot(losses.T[0], label='Discriminator') +plt.plot(losses.T[1], label='Generator') +plt.title("Training Losses") +plt.legend() +plt.show() + + +#Viewing the results of the GAN +def view_samples(epoch, samples): + + fig, axes = plt.subplots(figsize=(7,7), nrows=4, ncols=4, sharey=True, sharex=True) + fig.suptitle("Generated Digits") + for ax, img in zip(axes.flatten(), samples[epoch]): + img = img.detach().cpu().numpy() + ax.xaxis.set_visible(False) + ax.yaxis.set_visible(False) + im = ax.imshow(img.reshape((28,28)), cmap='Greys_r') + +with open('train_samples.pkl', 'rb') as f: + samples = pkl.load(f) + +view_samples(-1,samples) +plt.show()