Generative Adversarial Networks (GANs) have been a popular topic in the field of Deep Learning for generating synthetic data that looks like it was sampled from the training data. GANs have multiple applications such as image synthesis, data augmentation, image-to-image translation, and text generation.
In PyTorch, implementing and training GANs involves the following steps:
1. Importing necessary modules and setting up the device, such as GPU or CPU.
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision
import torchvision.transforms as transforms
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")2. Defining the discriminator and generator networks. Discriminator network tries to classify the input data into real or fake, while the generator network tries to generate fake data to fool the discriminator.
class Discriminator(nn.Module):
def __init__(self):
super(Discriminator, self).__init__()
self.conv_layer = nn.Sequential(
nn.Conv2d(in_channels=3, out_channels=64, kernel_size=4, stride=2),
nn.BatchNorm2d(64),
nn.LeakyReLU(0.2),
nn.Conv2d(in_channels=64, out_channels=128, kernel_size=4, stride=2),
nn.BatchNorm2d(128),
nn.LeakyReLU(0.2),
nn.Conv2d(in_channels=128, out_channels=256, kernel_size=4, stride=2),
nn.BatchNorm2d(256),
nn.LeakyReLU(0.2),
nn.Conv2d(in_channels=256, out_channels=512, kernel_size=4, stride=2),
nn.BatchNorm2d(512),
nn.LeakyReLU(0.2),
nn.Conv2d(in_channels=512, out_channels=1, kernel_size=4, stride=2),
nn.Sigmoid())
def forward(self, x):
# Forward pass
return self.conv_layer(x).view(-1, 1)
class Generator(nn.Module):
def __init__(self):
super(Generator, self).__init__()
self.deconv_layer = nn.Sequential(
nn.ConvTranspose2d(in_channels=100, out_channels=512, kernel_size=4, stride=1),
nn.BatchNorm2d(512),
nn.ReLU(True),
nn.ConvTranspose2d(in_channels=512, out_channels=256, kernel_size=4, stride=2),
nn.BatchNorm2d(256),
nn.ReLU(True),
nn.ConvTranspose2d(in_channels=256, out_channels=128, kernel_size=4, stride=2),
nn.BatchNorm2d(128),
nn.ReLU(True),
nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=4, stride=2),
nn.BatchNorm2d(64),
nn.ReLU(True),
nn.ConvTranspose2d(in_channels=64, out_channels=3, kernel_size=4, stride=2),
nn.Tanh())
def forward(self, x):
# Forward pass
return self.deconv_layer(x)3. Defining the loss functions and optimizers. The discriminator is trained to minimize binary cross-entropy loss, and the generator is trained to maximize the loss of the discriminator on fake images.
discriminator = Discriminator().to(device)
generator = Generator().to(device)
discriminator_criterion = nn.BCELoss()
generator_criterion = nn.BCELoss()
discriminator_optimizer = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))
generator_optimizer = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))4. Training the GAN model with real and fake data. Discriminator training occurs in one batch, while the generator training occurs in another batch.
for epoch in range(num_epochs):
for batch_idx, (real_images, _) in enumerate(trainloader):
real_images = real_images.to(device)
batch_size = real_images.shape[0]
# Train discriminator
discriminator.zero_grad()
real_targets = torch.ones(batch_size, 1).to(device)
fake_targets = torch.zeros(batch_size, 1).to(device)
# Real images
real_outputs = discriminator(real_images)
discriminator_loss_real = discriminator_criterion(real_outputs, real_targets)
# Fake images
noise = torch.randn(batch_size, 100, 1, 1).to(device)
fake_images = generator(noise)
fake_outputs = discriminator(fake_images.detach())
discriminator_loss_fake = discriminator_criterion(fake_outputs, fake_targets)
# Total loss
discriminator_loss = discriminator_loss_real + discriminator_loss_fake
discriminator_loss.backward()
discriminator_optimizer.step()
# Train generator
generator.zero_grad()
fake_outputs = discriminator(fake_images)
generator_loss = generator_criterion(fake_outputs, real_targets)
generator_loss.backward()
generator_optimizer.step()5. Measuring the performance of the GAN model. GANs can be challenging to measure performance, as no fixed loss function measures the quality of the generated samples. Some popular metrics that can be used for measuring the performance of GANs include the Frechet Inception Distance (FID) and Inception Score. FID measures the quality and diversity of generated images compared to the real data, while Inception Score measures the quality and the diversity of the generated images.
# Measuring FID
import torch.nn.functional as F
from torch.nn import Parameter
def compute_mu_sigma(x, eps=1e-6):
assert len(x.shape) == 4, "(N, C, H, W)"
N, C = x.shape[:2]
vals = x.view(N, C, -1)
mean = torch.mean(vals, dim=-1, keepdim=True)
centered = vals - mean
cov = torch.matmul(centered, centered.transpose(1, 2)) + (eps * torch.eye(C, device=x.device))
sigma = cov.sum(dim=-1) / (cov.shape[-1] - 1)
return mean.squeeze(), sigma
class GetFeatures(nn.Module):
def __init__(self):
super().__init__()
self.upsample = nn.Upsample(size=(299, 299), mode='bilinear', align_corners=True)
self.inception = torchvision.models.inception_v3(pretrained=True, aux_logits=False)
for param in self.inception.parameters():
param.requires_grad_(False)
def forward(self, x):
x = self.upsample(x)
x = self.inception(x)
return x
# Training loop...
# After training...
with torch.no_grad():
real_feature_list = []
fake_feature_list = []
get_features = GetFeatures().to(device)
for batch_idx, (real_images, _) in enumerate(trainloader):
real_images = real_images.to(device)
fake_images = generator(torch.randn(real_images.shape[0], 100, device=device))
# Get features for real images
real_features = get_features(real_images)
real_mu, real_sigma = compute_mu_sigma(real_features[:, :2048])
real_feature_list.append((real_mu.cpu().numpy(), real_sigma.cpu().numpy()))
# Get features for fake images
fake_features = get_features(fake_images)
fake_mu, fake_sigma = compute_mu_sigma(fake_features[:, :2048])
fake_feature_list.append((fake_mu.cpu().numpy(), fake_sigma.cpu().numpy()))
real_features = np.array(real_feature_list)
fake_features = np.array(fake_feature_list)
# Compute FID
m = ((real_features.mean(0) - fake_features.mean(0)) ** 2).sum()
s = np.linalg.norm(real_features.std(0) ** 2 + fake_features.std(0) ** 2 - 2 * (real_features.std(0) * fake_features.std(0)))**0.5
fid = m + s6. Techniques for stabilizing GAN training. GANs are known to suffer from training instability, including mode collapse and vanishing gradients. Some techniques that can be used to stabilize training include:
* Batch Normalization: normalize the activations to reduce internal covariance shift and to speed up the convergence during training.
* Weight Initialization: initialize the weights with low variance to avoid the vanishing gradients problem.
* Gradient Penalty: impose a penalty on the gradient norm of the discriminator to make sure that it does not saturate during training and that it does not reject generated samples due to sudden shifts in its behavior.
* Label Smoothing: replace the binary classification targets for the discriminator with a smooth target close to zero or one to reduce the discrimination between classes.
* WGAN-GP: Implementing the WGAN-GP architecture with the gradient penalty can be used to dramatically stabilize GANs.
Stabilizing GANs requires continuous experimentation with different architectures, hyperparameters, and loss functions to maintain a balance between discriminator and generator training.