GANs and VAEs

GAN Implementation

import torch
import torch.nn as nn

class Generator(nn.Module):
    def __init__(self, latent_dim=100, img_dim=784):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(latent_dim, 256),
            nn.LeakyReLU(0.2),
            nn.BatchNorm1d(256),
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2),
            nn.BatchNorm1d(512),
            nn.Linear(512, 1024),
            nn.LeakyReLU(0.2),
            nn.BatchNorm1d(1024),
            nn.Linear(1024, img_dim),
            nn.Tanh()
        )
    
    def forward(self, z):
        return self.model(z)

class Discriminator(nn.Module):
    def __init__(self, img_dim=784):
        super().__init__()
        self.model = nn.Sequential(
            nn.Linear(img_dim, 512),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2),
            nn.Dropout(0.3),
            nn.Linear(256, 1),
            nn.Sigmoid()
        )
    
    def forward(self, img):
        return self.model(img)

class GAN:
    def __init__(self, latent_dim=100, img_dim=784):
        self.generator = Generator(latent_dim, img_dim)
        self.discriminator = Discriminator(img_dim)
        self.latent_dim = latent_dim
        
        self.g_optimizer = torch.optim.Adam(self.generator.parameters(), lr=0.0002, betas=(0.5, 0.999))
        self.d_optimizer = torch.optim.Adam(self.discriminator.parameters(), lr=0.0002, betas=(0.5, 0.999))
        self.criterion = nn.BCELoss()
    
    def train_step(self, real_images):
        batch_size = real_images.shape[0]
        
        real_labels = torch.ones(batch_size, 1)
        fake_labels = torch.zeros(batch_size, 1)
        
        real_output = self.discriminator(real_images)
        d_loss_real = self.criterion(real_output, real_labels)
        
        z = torch.randn(batch_size, self.latent_dim)
        fake_images = self.generator(z)
        fake_output = self.discriminator(fake_images.detach())
        d_loss_fake = self.criterion(fake_output, fake_labels)
        
        d_loss = d_loss_real + d_loss_fake
        self.d_optimizer.zero_grad()
        d_loss.backward()
        self.d_optimizer.step()
        
        z = torch.randn(batch_size, self.latent_dim)
        fake_images = self.generator(z)
        fake_output = self.discriminator(fake_images)
        g_loss = self.criterion(fake_output, real_labels)
        
        self.g_optimizer.zero_grad()
        g_loss.backward()
        self.g_optimizer.step()
        
        return d_loss.item(), g_loss.item()
    
    def generate(self, n_samples):
        z = torch.randn(n_samples, self.latent_dim)
        return self.generator(z)

gan = GAN(latent_dim=100)
for epoch in range(100):
    d_loss, g_loss = gan.train_step(real_data)

VAE Implementation

class VAE(nn.Module):
    def __init__(self, input_dim=784, hidden_dim=400, latent_dim=20):
        super().__init__()
        
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU()
        )
        
        self.mu_layer = nn.Linear(hidden_dim, latent_dim)
        self.logvar_layer = nn.Linear(hidden_dim, latent_dim)
        
        self.decoder = nn.Sequential(
            nn.Linear(latent_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.ReLU(),
            nn.Linear(hidden_dim, input_dim),
            nn.Sigmoid()
        )
    
    def encode(self, x):
        h = self.encoder(x)
        return self.mu_layer(h), self.logvar_layer(h)
    
    def reparameterize(self, mu, logvar):
        std = torch.exp(0.5 * logvar)
        eps = torch.randn_like(std)
        return mu + eps * std
    
    def decode(self, z):
        return self.decoder(z)
    
    def forward(self, x):
        mu, logvar = self.encode(x)
        z = self.reparameterize(mu, logvar)
        recon = self.decode(z)
        return recon, mu, logvar
    
    def loss_function(self, recon_x, x, mu, logvar):
        recon_loss = nn.functional.binary_cross_entropy(recon_x, x, reduction='sum')
        
        kl_loss = -0.5 * torch.sum(1 + logvar - mu.pow(2) - logvar.exp())
        
        return recon_loss + kl_loss

vae = VAE(input_dim=784, latent_dim=20)
optimizer = torch.optim.Adam(vae.parameters(), lr=1e-3)

for epoch in range(50):
    for batch in dataloader:
        recon, mu, logvar = vae(batch)
        loss = vae.loss_function(recon, batch, mu, logvar)
        
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

Best Practices

• Use spectral normalization in GANs for stability
• Apply label smoothing for discriminator training
• Use latent space interpolation for smooth transitions
• Monitor mode coverage in GANs
• Use KL annealing in VAEs for better training
• Consider VQ-VAE for discrete representations