
import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, transforms
import matplotlib.pyplot as plt

# Set the computation device: GPU if available, else CPU.
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# Generator model definition
class Generator(nn.Module):
    def __init__(self, latent_dim):
        super(Generator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(latent_dim, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, 256),
            nn.BatchNorm1d(256, 0.8),
            nn.ReLU(inplace=True),
            nn.Linear(256, 512),
            nn.BatchNorm1d(512, 0.8),
            nn.ReLU(inplace=True),
            nn.Linear(512, 28 * 28),
            nn.Tanh()
        )
    
    def forward(self, z):
        img = self.model(z)
        img = img.view(img.size(0), 1, 28, 28)
        return img

# Discriminator model definition
class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator, self).__init__()
        self.model = nn.Sequential(
            nn.Linear(28 * 28, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )
    
    def forward(self, img):
        # Flatten the image from (batch, 1, 28, 28) to (batch, 784)
        img_flat = img.view(img.size(0), -1)
        validity = self.model(img_flat)
        return validity

# Hyperparameters
latent_dim = 100
batch_size = 64
epochs = 300

# Initialize generator and discriminator and move to device.
generator = Generator(latent_dim).to(device)
discriminator = Discriminator().to(device)

# Optimizers and loss criterion
optimizer_G = optim.Adam(generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
optimizer_D = optim.Adam(discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))
adversarial_loss = nn.BCELoss()

# DataLoader for the MNIST dataset.
dataloader = torch.utils.data.DataLoader(
    datasets.MNIST(
        "./data", train=True, download=True,
        transform=transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.5,), (0.5,))
        ])
    ),
    batch_size=batch_size, shuffle=True
)

# For recording average loss per epoch.
epoch_d_losses = []
epoch_g_losses = []

# This list will store the 5-image samples from every 20th epoch.
sample_images_rows = []


# Training loop.
for epoch in range(epochs):
    epoch_d_loss = 0.0
    epoch_g_loss = 0.0
    for i, (imgs, _) in enumerate(dataloader):
        # Transfer images to device.
        imgs = imgs.to(device)
        valid = torch.ones((imgs.size(0), 1), device=device)
        fake = torch.zeros((imgs.size(0), 1), device=device)
        
        # Train Generator.
        optimizer_G.zero_grad()
        z = torch.randn(imgs.size(0), latent_dim, device=device)
        gen_imgs = generator(z)
        g_loss = adversarial_loss(discriminator(gen_imgs), valid)
        g_loss.backward()
        optimizer_G.step()
        
        # Train Discriminator.
        optimizer_D.zero_grad()
        real_loss = adversarial_loss(discriminator(imgs), valid)
        fake_loss = adversarial_loss(discriminator(gen_imgs.detach()), fake) ## detach is important so we don't mess with the gradient in the generator
        d_loss = (real_loss + fake_loss) / 2
        d_loss.backward()
        optimizer_D.step()
        
        epoch_g_loss += g_loss.item()
        epoch_d_loss += d_loss.item()

        if i % 100 == 0:
            print(f"[Epoch {epoch+1}/{epochs}] [Batch {i}/{len(dataloader)}] "
                  f"[D loss: {d_loss.item():.4f}] [G loss: {g_loss.item():.4f}]")
    
    # Average losses for the epoch.
    avg_d_loss = epoch_d_loss / len(dataloader)
    avg_g_loss = epoch_g_loss / len(dataloader)
    epoch_d_losses.append(avg_d_loss)
    epoch_g_losses.append(avg_g_loss)
    
    # Every 20th epoch, generate and store a sample of 5 images.
    if (epoch + 1) % 20 == 0:
        with torch.no_grad():
            sample_z = torch.randn(5, latent_dim, device=device)
            sample_imgs = generator(sample_z).detach().cpu()
        sample_images_rows.append(sample_imgs)
        print("Saved sample images for epoch", epoch+1)


# Plot the average generator and discriminator losses over epochs.
plt.figure(figsize=(10, 5))
plt.plot(range(1, epochs+1), epoch_d_losses, label="Discriminator Loss", marker='o')
plt.plot(range(1, epochs+1), epoch_g_losses, label="Generator Loss", marker='o')
plt.title("Average Generator and Discriminator Loss per Epoch")
plt.xlabel("Epoch")
plt.ylabel("Loss")
plt.legend()
plt.grid(True)
plt.show()

# Combine all sample rows (from every 20th epoch) into one grid.
num_samples = len(sample_images_rows)  # Should be 15 for 300 epochs.
fig, axes = plt.subplots(num_samples, 5, figsize=(15, num_samples * 3))
for row_idx, sample_imgs in enumerate(sample_images_rows):
    for col_idx in range(5):
        axes[row_idx, col_idx].imshow(sample_imgs[col_idx].view(28, 28), cmap='gray')
        axes[row_idx, col_idx].axis('off')
plt.suptitle("Sample Images Every 20th Epoch", fontsize=16)
plt.tight_layout()
plt.show()

Lecture 19 - Generative adversarial networks¶

ECE364 - Programming Methods for Machine Learning¶

Nickvash Kani¶

Slides based off prior lectures by Alex Schwing, Aigou Han, Farzad Kamalabadi, Corey Snyder. All mistakes are my own!¶

Unsupervised learning (continued)¶

Implicit generative models¶

Visualizing the image space¶

Generative Adversarial Networks¶

Generative Adversarial Networks (GANs)¶

GAN loss function¶

Modified loss function¶

Example: GAN on MNIST¶

How to Train a GAN?¶

Results from above code:¶

GPU training¶

Assuming you have a model and a DataLoader¶

Other uses for GANs¶

Generating celebrity faces [4]¶

Style transfer with CycleGan [5]¶

Text to image with StackGan [7]¶

Other examples for GANs [2]¶

That's it for today¶

References¶