Creating Faces from Scratch: A Deep Dive into GANs with PyTorch

import torch
import torch.nn as nn
import torchvision.datasets as dset
from torchvision import transforms
import torchvision.utils as vutils
import torch.optim as optim

import matplotlib.pyplot as plt
import numpy as np
import matplotlib.animation as animation
from IPython.display import HTML

IMAGE_SIZE = 64

BATCH_SIZE = 512
WORKERS = 4

NGPUS = 2

NUM_CHANNELS = 3
LATENT_DIM = 100

G_FEATURES = 64
D_FEATURES = 64

BETA1=0.5
LR=0.0002
NUM_EPOCHS=10

transform = transforms.Compose([
    transforms.Resize(IMAGE_SIZE),
    transforms.CenterCrop(IMAGE_SIZE),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
])

dataset = dset.ImageFolder('./data/celeba', transform=transform)

dataloader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE,
                                         shuffle=True, num_workers=WORKERS)

device = torch.device("cuda:0" if (torch.cuda.is_available() and NGPUS > 0) else "cpu")

# Plot some training images
real_batch = next(iter(dataloader))
plt.figure(figsize=(8,8))
plt.axis("off")
plt.title("Training Images")
plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=2, normalize=True).cpu(),(1,2,0)))
plt.show()

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)

class Generator(nn.Module):

    def __init__(self):
        super(Generator, self).__init__()
        self.model = nn.Sequential(
            # input is Z_DIM, going into a convolutional layer
            nn.ConvTranspose2d(LATENT_DIM, G_FEATURES * 8, 4, 1, 0),
            nn.BatchNorm2d(G_FEATURES * 8),
            nn.ReLU(True),

            # state size. (G_FEATURES * 8) x 4 x 4
            nn.ConvTranspose2d(G_FEATURES * 8, G_FEATURES * 4, 4, 2, 1),
            nn.BatchNorm2d(G_FEATURES * 4),
            nn.ReLU(True),

            # state size. (G_FEATURES * 4) x 8 x 8
            nn.ConvTranspose2d(G_FEATURES * 4, G_FEATURES * 2, 4, 2, 1),
            nn.BatchNorm2d(G_FEATURES * 2),
            nn.ReLU(True),

            # state size. (G_FEATURES * 2) x 16 x 16
            nn.ConvTranspose2d(G_FEATURES * 2, G_FEATURES, 4, 2, 1),
            nn.BatchNorm2d(G_FEATURES),
            nn.ReLU(True),

            # state size. (G_FEATURES) x 32 x 32
            nn.ConvTranspose2d(G_FEATURES, NUM_CHANNELS, 4, 2, 1),
            nn.Tanh()
        )
    def forward(self, z):
        return self.model(z)

netG = Generator().to(device)

if (device.type == 'cuda') and (NGPUS > 1):
    netG = nn.DataParallel(netG, list(range(NGPUS)))

netG.apply(weights_init)

netG

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        
        self.model = nn.Sequential(
            nn.Conv2d(NUM_CHANNELS, D_FEATURES, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(D_FEATURES, D_FEATURES * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(D_FEATURES * 2),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(D_FEATURES * 2, D_FEATURES * 4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(D_FEATURES * 4),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(D_FEATURES * 4, D_FEATURES * 8, 4, 2, 1, bias=False),
            nn.BatchNorm2d(D_FEATURES * 8),
            nn.LeakyReLU(0.2, inplace=True),

            nn.Conv2d(D_FEATURES * 8, 1, 4, 1, 0, bias=False),
            nn.Sigmoid(),
        )

    def forward(self, x):
        return self.model(x)

netD = Discriminator().to(device)

if (device.type == 'cuda') and (NGPUS > 1):
    netD = nn.DataParallel(netD, list(range(NGPUS)))

netD.apply(weights_init)

netD

criterion = nn.BCELoss()

fixed_noise = torch.randn(64, LATENT_DIM, 1, 1, device=device)

real_label = 1.
fake_label = 0.

optimizerD = optim.Adam(netD.parameters(), lr=LR, betas=(BETA1, 0.999))
optimizerG = optim.Adam(netG.parameters(), lr=LR, betas=(BETA1, 0.999))

# Initialize tracking variables for training progress
img_list = []    # Store generated images at checkpoints
G_losses = []    # Track generator losses over time
D_losses = []    # Track discriminator losses over time
iters = 0        # Global iteration counter

# DCGAN Training Loop
# Reference: https://docs.pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html#training
print("Starting Training Loop...")

for epoch in range(NUM_EPOCHS):
    for i, data in enumerate(dataloader):
        
        ############################
        # PHASE 1: Update Discriminator Network
        # Goal: maximize log(D(x)) + log(1 - D(G(z)))
        # This means: correctly classify real images as real AND fake images as fake
        ############################
        
        # === TRAIN DISCRIMINATOR ON REAL IMAGES ===
        netD.zero_grad()  # Clear discriminator gradients
        
        # Prepare real image batch
        real_cpu = data[0].to(device)
        b_size = real_cpu.size(0)
        label = torch.full((b_size,), real_label, dtype=torch.float, device=device)
        
        # Forward pass: discriminator classifies real images
        output = netD(real_cpu).view(-1)
        
        # Calculate loss: how well D identifies real images as real
        errD_real = criterion(output, label)
        
        # Backward pass: compute gradients for real image classification
        errD_real.backward()
        D_x = output.mean().item()  # Average discriminator output on real images
        
        
        # === TRAIN DISCRIMINATOR ON FAKE IMAGES ===
        # Generate batch of random noise vectors
        noise = torch.randn(b_size, LATENT_DIM, 1, 1, device=device)
        
        # Generate fake images using generator
        fake = netG(noise)
        label.fill_(fake_label)  # Set labels to "fake" for this batch
        
        # Forward pass: discriminator classifies fake images
        # Note: .detach() prevents gradients from flowing back to generator
        output = netD(fake.detach()).view(-1)
        
        # Calculate loss: how well D identifies fake images as fake
        errD_fake = criterion(output, label)
        
        # Backward pass: accumulate gradients with previous real image gradients
        errD_fake.backward()
        D_G_z1 = output.mean().item()  # Average discriminator output on fake images
        
        # Total discriminator error and update
        errD = errD_real + errD_fake
        optimizerD.step()  # Apply discriminator parameter updates
        
        
        ############################
        # PHASE 2: Update Generator Network
        # Goal: maximize log(D(G(z)))
        # This means: fool the discriminator into thinking fake images are real
        ############################
        
        netG.zero_grad()  # Clear generator gradients
        
        # We want the discriminator to classify our fake images as real
        label.fill_(real_label)
        
        # Forward pass: discriminator re-evaluates the same fake images
        # Note: No .detach() here - we want gradients to flow back to generator
        output = netD(fake).view(-1)
        
        # Calculate generator loss: how well G fools the discriminator
        errG = criterion(output, label)
        
        # Backward pass: compute gradients for generator
        errG.backward()
        D_G_z2 = output.mean().item()  # Discriminator output on fake images (after D update)
        
        # Update generator parameters
        optimizerG.step()
        
        
        ############################
        # LOGGING AND CHECKPOINTS
        ############################
        
        # Print training statistics every 50 iterations
        if i % 50 == 0:
            print('[%d/%d][%d/%d]\tLoss_D: %.4f\tLoss_G: %.4f\tD(x): %.4f\tD(G(z)): %.4f / %.4f'
                  % (epoch, NUM_EPOCHS, i, len(dataloader),
                     errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
        
        # Save losses for plotting training curves later
        G_losses.append(errG.item())
        D_losses.append(errD.item())
        
        # Generate and save sample images at checkpoints
        if (iters % 500 == 0) or ((epoch == NUM_EPOCHS-1) and (i == len(dataloader)-1)):
            with torch.no_grad():  # Disable gradient computation for inference
                fake = netG(fixed_noise).detach().cpu()  # Generate images from fixed noise
            # Create grid of images and add to list for visualization
            img_list.append(vutils.make_grid(fake, padding=2, normalize=True))
        
        iters += 1  # Increment global iteration counter

plt.figure(figsize=(10,5))
plt.title("Generator and Discriminator Loss During Training")
plt.plot(G_losses,label="G")
plt.plot(D_losses,label="D")
plt.xlabel("iterations")
plt.ylabel("Loss")
plt.legend()
plt.show()

fig = plt.figure(figsize=(8,8))
plt.axis("off")
ims = [[plt.imshow(np.transpose(i,(1,2,0)), animated=True)] for i in img_list]
ani = animation.ArtistAnimation(fig, ims, interval=1000, repeat_delay=1000, blit=True)

HTML(ani.to_jshtml())

# Grab a batch of real images from the dataloader
real_batch = next(iter(dataloader))

# Plot the real images
plt.figure(figsize=(15,15))
plt.subplot(1,2,1)
plt.axis("off")
plt.title("Real Images")
plt.imshow(np.transpose(vutils.make_grid(real_batch[0].to(device)[:64], padding=5, normalize=True).cpu(),(1,2,0)))

# Plot the fake images from the last epoch
plt.subplot(1,2,2)
plt.axis("off")
plt.title("Fake Images")
plt.imshow(np.transpose(img_list[-1],(1,2,0)))
plt.show()