Creating New Anime Faces with Generative Adversarial Networks (GANs) using PyTorch

Generative adversarial networks (GANs), in simple terms, are two neural network players competing against each other for the sole purpose of defeating the other. In this process, both of the neural networks learn a lot and develop more accurate results which can be used by us to create new, synthetic data. This new data is very close in resemblance to the original data used by the networks.

GANs are being used in image generation, voice generation and video generation. They fall under Generative Modeling which is an unsupervised learning task in ML involved in generating new content.

The generator generates a “fake” sample given a random vector/matrix, and the discriminator attempts to detect whether a given sample is “real” (picked from the training data) or “fake” (generated by the generator). Training happens in tandem: we train the discriminator for a few epochs, then train the generator for a few epochs, and repeat. This way both the generator and the discriminator get better at doing their jobs. GANs however, can be notoriously difficult to train, and are extremely sensitive to hyperparameters, activation functions and regularization.

In this tutorial, we’ll train a GAN to generate images of anime characters’ faces. We’ll use the Anime Face Dataset, which consists of over 63,000 cropped anime faces. There are no labels here as this is an unsupervised learning technique.

I have used Google Colab for coding the problem. So, I will show you how it is done step-wise.

Downloading the Data

We will download the data with the ‘opendatasets’ library. The dataset is located at ‘https://www.kaggle.com/splcher/animefacedataset’.

import opendatasets as od
dataset_url = ‘https://www.kaggle.com/splcher/animefacedataset'
od.download(dataset_url)

You will need to have a Kaggle account for downloading this dataset as the above command will ask for your Kaggle account ID and the API key. You can find the API key in the account details of your Kaggle account. Once you enter those, your dataset will be downloaded in Colab. The dataset has a single folder called images which contains all 63,000+ images in JPG format.

Exploring the Data

import os
DATA_DIR = ‘./animefacedataset’
print(os.listdir(DATA_DIR))
print(os.listdir(DATA_DIR+’/images’)[:10])

You get the output as [‘14791_2006.jpg’, ‘15606_2006.jpg’, ‘37864_2012.jpg’, ‘61376_2018.jpg’, ‘51585_2015.jpg’, ‘1505_2001.jpg’, ‘42640_2013.jpg’, ‘54069_2016.jpg’, ‘48489_2014.jpg’, ‘388_2000.jpg’].

Importing Required Libraries

from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder
import torchvision.transforms as T
import torch
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
import torch.nn as nn
%matplotlib inline
from torchvision.utils import save_image
from tqdm.notebook import tqdm
import torch.nn.functional as F
from IPython.display import Image

Creating the DataLoader

Let’s load this dataset using the ImageFolder class from torchvision. We will also resize and crop the images to 64x64 px, and normalize the pixel values with a mean & standard deviation of 0.5 for each channel. This will ensure that pixel values are in the range (-1, 1), which is more convenient for training the discriminator. We will also create a data loader to load the data in batches.

image_size = 64
batch_size = 128
stats = (0.5, 0.5, 0.5), (0.5, 0.5, 0.5)
train_ds = ImageFolder(DATA_DIR, transform=T.Compose([
T.Resize(image_size),
T.CenterCrop(image_size),
T.ToTensor(),
T.Normalize(*stats)]))
train_dl = DataLoader(train_ds, batch_size, shuffle=True, num_workers=3, pin_memory=True)

Creating Helper Functions for our Work

Function for De-normalization:

def denorm(img_tensors):
return img_tensors * stats[1][0] + stats[0][0]

Function for Showing the Images:

def show_images(images, nmax=64):
fig, ax = plt.subplots(figsize=(8, 8))
ax.set_xticks([]); ax.set_yticks([])
ax.imshow(make_grid(denorm(images.detach()[:nmax]), nrow=8).permute(1, 2, 0))

Function for Showing the Batch:

def show_batch(dl, nmax=64):
for images, _ in dl:
show_images(images, nmax)
break

To check out the batch, we do,

show_batch(train_dl)

We can see that the batch looks like below.

Setting up the GPU

Below are the functions to be used for setting the GPU.

def get_default_device():
“””Pick GPU if available, else CPU”””
if torch.cuda.is_available():
return torch.device(‘cuda’)
else:
return torch.device(‘cpu’)
def to_device(data, device):
“””Move tensor(s) to chosen device”””
if isinstance(data, (list,tuple)):
return [to_device(x, device) for x in data]
return data.to(device, non_blocking=True)
class DeviceDataLoader():
“””Wrap a dataloader to move data to a device”””
def __init__(self, dl, device):
self.dl = dl
self.device = device
def __iter__(self):
“””Yield a batch of data after moving it to device”””
for b in self.dl:
yield to_device(b, self.device)
def __len__(self):
“””Number of batches”””
return len(self.dl)
Based on where you’re running this notebook, your default device could be a CPU (torch.device('cpu')) or a GPU (torch.device('cuda'))
device = get_default_device()
train_dl = DeviceDataLoader(train_dl, device)

After this, we are now able to use the GPU.

Creating the Discriminator

The discriminator takes an image as input, and tries to classify it as “real” or “generated”. In this sense, it’s like any other neural network. We’ll use a convolutional neural networks (CNN) which outputs a single number output for every image. We’ll use stride of 2 to progressively reduce the size of the output feature map.

We create the discriminator as follows.

discriminator = nn.Sequential(
# in: 3 x 64 x 64
nn.Conv2d(3, 64, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2, inplace=True),
# out: 64 x 32 x 32
nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2, inplace=True),
# out: 128 x 16 x 16
nn.Conv2d(128, 256, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2, inplace=True),
# out: 256 x 8 x 8
nn.Conv2d(256, 512, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.2, inplace=True),
# out: 512 x 4 x 4
nn.Conv2d(512, 1, kernel_size=4, stride=1, padding=0, bias=False),
# out: 1 x 1 x 1
nn.Flatten(),
nn.Sigmoid())

We’re using the Leaky ReLU activation for the discriminator. Different from the regular ReLU function, Leaky ReLU allows the pass of a small gradient signal for negative values. As a result, it makes the gradients from the discriminator flows stronger into the generator. Instead of passing a gradient (slope) of 0 in the back-prop pass, it passes a small negative gradient. Just like any other binary classification model, the output of the discriminator is a single number between 0 and 1, which can be interpreted as the probability of the input image being real i.e. picked from the original dataset.

Let’s move the discriminator model to the chosen device.

discriminator = to_device(discriminator, device)

Creating the Generator

The input to the generator is typically a vector or a matrix of random numbers (referred to as a latent tensor) which is used as a seed for generating an image. The generator will convert a latent tensor of shape (128, 1, 1) into an image tensor of shape 3 x 28 x 28. To achive this, we'll use the ConvTranspose2d layer from PyTorch.

latent_size = 128
generator = nn.Sequential(
# in: latent_size x 1 x 1
nn.ConvTranspose2d(latent_size, 512, kernel_size=4, stride=1, padding=0, bias=False),
nn.BatchNorm2d(512),
nn.ReLU(True),
# out: 512 x 4 x 4
nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(256),
nn.ReLU(True),
# out: 256 x 8 x 8
nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(128),
nn.ReLU(True),
# out: 128 x 16 x 16
nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1, bias=False),
nn.BatchNorm2d(64),
nn.ReLU(True),
# out: 64 x 32 x 32
nn.ConvTranspose2d(64, 3, kernel_size=4, stride=2, padding=1, bias=False),
nn.Tanh()
# out: 3 x 64 x 64
)

The ReLU activation is used in the generator with the exception of the output layer which uses the Tanh function. We observed that using a bounded activation allowed the model to learn more quickly to saturate and cover the color space of the training distribution. Within the discriminator we found the leaky rectified activation to work well, especially for higher resolution modeling.

Note that since the outputs of the TanH activation lie in the range [-1,1], we have applied the similar transformation to the images in the training dataset. Let's generate some outputs using the generator and view them as images by transforming and denormalizing the output.

xb = torch.randn(batch_size, latent_size, 1, 1) # random latent tensors
fake_images = generator(xb)
print(fake_images.shape)
show_images(fake_images)

As one might expect, the output from the generator is basically random noise, since we haven’t trained it yet.

Let’s move the generator to the chosen device.

generator = to_device(generator, device)

Training the Discriminator

Since the discriminator is a binary classification model, we can use the binary cross entropy loss function to quantify how well it is able to differentiate between real and generated images.

def train_discriminator(real_images, opt_d):
# Clear discriminator gradients
opt_d.zero_grad()
# Pass real images through discriminator
real_preds = discriminator(real_images)
real_targets = torch.ones(real_images.size(0), 1, device=device)
real_loss = F.binary_cross_entropy(real_preds, real_targets)
real_score = torch.mean(real_preds).item()
# Generate fake images
latent = torch.randn(batch_size, latent_size, 1, 1, device=device)
fake_images = generator(latent)
# Pass fake images through discriminator
fake_targets = torch.zeros(fake_images.size(0), 1, device=device)
fake_preds = discriminator(fake_images)
fake_loss = F.binary_cross_entropy(fake_preds, fake_targets)
fake_score = torch.mean(fake_preds).item()
# Update discriminator weights
loss = real_loss + fake_loss
loss.backward()
opt_d.step()
return loss.item(), real_score, fake_score

Here are the steps involved in training the discriminator.

• We expect the discriminator to output 1 if the image was picked from the real MNIST dataset, and 0 if it was generated using the generator network.
• We first pass a batch of real images, and compute the loss, setting the target labels to 1.
• Then we pass a batch of fake images (generated using the generator) pass them into the discriminator, and compute the loss, setting the target labels to 0.
• Finally we add the two losses and use the overall loss to perform gradient descent to adjust the weights of the discriminator.

It’s important to note that we don’t change the weights of the generator model while training the discriminator (opt_d only affects the discriminator.parameters())

Training the Generator

Since the outputs of the generator are images, it’s not obvious how we can train the generator. This is where we employ a rather elegant trick, which is to use the discriminator as a part of the loss function. Here’s how it works:

• We generate a batch of images using the generator, pass the into the discriminator.
• We calculate the loss by setting the target labels to 1 i.e. real. We do this because the generator’s objective is to “fool” the discriminator.
• We use the loss to perform gradient descent i.e. change the weights of the generator, so it gets better at generating real-like images to “fool” the discriminator.

Here’s what this looks like in code.

def train_generator(opt_g):
# Clear generator gradients
opt_g.zero_grad()
# Generate fake images
latent = torch.randn(batch_size, latent_size, 1, 1, device=device)
fake_images = generator(latent)
# Try to fool the discriminator
preds = discriminator(fake_images)
targets = torch.ones(batch_size, 1, device=device)
loss = F.binary_cross_entropy(preds, targets)
# Update generator weights
loss.backward()
opt_g.step()
return loss.item()

Now, let’s create a directory to save the images.

sample_dir = ‘generated’
os.makedirs(sample_dir, exist_ok=True)

Let’s also create a helper function to export the generated images.

def save_samples(index, latent_tensors, show=True):
fake_images = generator(latent_tensors)
fake_fname = ‘generated-images-{0:0=4d}.png’.format(index)
save_image(denorm(fake_images), os.path.join(sample_dir, fake_fname), nrow=8)
print(‘Saving’, fake_fname)
if show:
fig, ax = plt.subplots(figsize=(8, 8))
ax.set_xticks([]); ax.set_yticks([])
ax.imshow(make_grid(fake_images.cpu().detach(), nrow=8).permute(1, 2, 0))

We’ll use a fixed set of input vectors to the generator to see how the individual generated images evolve over time as we train the model. Let’s save one set of images before we start training our model.

fixed_latent = torch.randn(64, latent_size, 1, 1, device=device)
save_samples(0, fixed_latent)

The generated images look like below.

Let’s define a fit function to train the discriminator and generator in tandem for each batch of training data. We'll use the Adam optimizer with some custom parameters (betas) that are known to work well for GANs. We will also save some sample generated images at regular intervals for inspection.

def fit(epochs, lr, start_idx=1):
torch.cuda.empty_cache()
# Losses & scores
losses_g = []
losses_d = []
real_scores = []
fake_scores = []
# Create optimizers
opt_d = torch.optim.Adam(discriminator.parameters(), lr=lr, betas=(0.5, 0.999))
opt_g = torch.optim.Adam(generator.parameters(), lr=lr, betas=(0.5, 0.999))
for epoch in range(epochs):
for real_images, _ in tqdm(train_dl):
# Train discriminator
loss_d, real_score, fake_score = train_discriminator(real_images, opt_d)
# Train generator
loss_g = train_generator(opt_g)
# Record losses & scores
losses_g.append(loss_g)
losses_d.append(loss_d)
real_scores.append(real_score)
fake_scores.append(fake_score)
# Log losses & scores (last batch)
print(“Epoch [{}/{}], loss_g: {:.4f}, loss_d: {:.4f}, real_score: {:.4f}, fake_score: {:.4f}”.format(
epoch+1, epochs, loss_g, loss_d, real_score, fake_score))
# Save generated images
save_samples(epoch+start_idx, fixed_latent, show=False)
return losses_g, losses_d, real_scores, fake_scores

Training the Model

lr = 0.0002
epochs = 25

We choose learning rate 0.0002 and number of epochs as 25 for the model. But, you can choose any combination and see how the model works for yourself.

history = fit(epochs, lr)

At the last epoch, we get,

Epoch [25/25], loss_g: 2.4400, loss_d: 0.1189, real_score: 0.9247, fake_score: 0.0280 Saving generated-images-0025.png

Saving the Model

torch.save(generator.state_dict(), ‘G.pth’)
torch.save(discriminator.state_dict(), ‘D.pth’)

Viewing the Progress

As we see the model gets better and in the end, we obtain previously unseen anime faces that look like the real thing!

Plotting the Graphs

plt.plot(losses_d, ‘-’)
plt.plot(losses_g, ‘-’)
plt.xlabel(‘epoch’)
plt.ylabel(‘loss’)
plt.legend([‘Discriminator’, ‘Generator’])
plt.title(‘Losses’);

plt.plot(real_scores, ‘-’)
plt.plot(fake_scores, ‘-’)
plt.xlabel(‘epoch’)
plt.ylabel(‘score’)
plt.legend([‘Real’, ‘Fake’])
plt.title(‘Scores’);

So, this is it! Hope this article provides you the motivation to try out different ideas with GANs.