# 📚 PyTorch Tutorial for Beginners - Part 5/6: Generative Models & Advanced Topics
#PyTorch #DeepLearning #GANs #VAEs #ReinforcementLearning #Deployment

Welcome to Part 5 of our PyTorch series! This comprehensive lesson explores generative modeling, reinforcement learning, model optimization, and deployment strategies with practical implementations.

---

## 🔹 Generative Adversarial Networks (GANs)
### 1. GAN Core Concepts

Key Components:
- Generator: Creates fake samples from noise (typically a transposed CNN)
- Discriminator: Distinguishes real vs. fake samples (CNN classifier)
- Adversarial Training: The two networks compete in a minimax game

### 2. DCGAN Implementation

class Generator(nn.Module):
    def __init__(self, latent_dim, img_channels, features_g):
        super().__init__()
        self.net = nn.Sequential(
            # Input: N x latent_dim x 1 x 1
            nn.ConvTranspose2d(latent_dim, features_g*8, 4, 1, 0, bias=False),
            nn.BatchNorm2d(features_g*8),
            nn.ReLU(),
            # 4x4
            nn.ConvTranspose2d(features_g*8, features_g*4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features_g*4),
            nn.ReLU(),
            # 8x8
            nn.ConvTranspose2d(features_g*4, features_g*2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features_g*2),
            nn.ReLU(),
            # 16x16
            nn.ConvTranspose2d(features_g*2, img_channels, 4, 2, 1, bias=False),
            nn.Tanh()
            # 32x32
        )

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

class Discriminator(nn.Module):
    def __init__(self, img_channels, features_d):
        super().__init__()
        self.net = nn.Sequential(
            # Input: N x img_channels x 32 x 32
            nn.Conv2d(img_channels, features_d, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2),
            # 16x16
            nn.Conv2d(features_d, features_d*2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features_d*2),
            nn.LeakyReLU(0.2),
            # 8x8
            nn.Conv2d(features_d*2, features_d*4, 4, 2, 1, bias=False),
            nn.BatchNorm2d(features_d*4),
            nn.LeakyReLU(0.2),
            # 4x4
            nn.Conv2d(features_d*4, 1, 4, 1, 0, bias=False),
            nn.Sigmoid()
        )

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

# Initialize
gen = Generator(latent_dim=100, img_channels=3, features_g=64).to(device)
disc = Discriminator(img_channels=3, features_d=64).to(device)

# Loss and optimizers
criterion = nn.BCELoss()
opt_gen = optim.Adam(gen.parameters(), lr=0.0002, betas=(0.5, 0.999))
opt_disc = optim.Adam(disc.parameters(), lr=0.0002, betas=(0.5, 0.999))

### 3. GAN Training Loop

def train_gan(gen, disc, loader, num_epochs):
    fixed_noise = torch.randn(32, 100, 1, 1).to(device)
    
    for epoch in range(num_epochs):
        for batch_idx, (real, _) in enumerate(loader):
            real = real.to(device)
            noise = torch.randn(real.size(0), 100, 1, 1).to(device)
            fake = gen(noise)
            
            # Train Discriminator
            disc_real = disc(real).view(-1)
            loss_disc_real = criterion(disc_real, torch.ones_like(disc_real))
            disc_fake = disc(fake.detach()).view(-1)
            loss_disc_fake = criterion(disc_fake, torch.zeros_like(disc_fake))
            loss_disc = (loss_disc_real + loss_disc_fake) / 2
            disc.zero_grad()
            loss_disc.backward()
            opt_disc.step()
            
            # Train Generator
            output = disc(fake).view(-1)
            loss_gen = criterion(output, torch.ones_like(output))
            gen.zero_grad()
            loss_gen.backward()
            opt_gen.step()
            
        # Visualization
        with torch.no_grad():
            fake = gen(fixed_noise)
            save_image(fake, f"gan_samples/epoch_{epoch}.png", normalize=True)

### 2. Pruning

parameters_to_prune = (
    (model.conv1, 'weight'),
    (model.fc1, 'weight'),
)

prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.L1Unstructured,
    amount=0.2
)

# Remove pruning reparameterization
for module, param in parameters_to_prune:
    prune.remove(module, param)

### 3. ONNX Export

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    }
)

### 4. TorchScript

# Tracing
example_input = torch.rand(1, 3, 224, 224)
traced_script = torch.jit.trace(model, example_input)
traced_script.save("traced_model.pt")

# Scripting
scripted_model = torch.jit.script(model)
scripted_model.save("scripted_model.pt")

---

## 🔹 PyTorch Lightning Best Practices
### 1. LightningModule Structure

import pytorch_lightning as pl

class LitModel(pl.LightningModule):
    def __init__(self, learning_rate=1e-3):
        super().__init__()
        self.save_hyperparameters()
        self.model = nn.Sequential(
            nn.Linear(28*28, 128),
            nn.ReLU(),
            nn.Linear(128, 10)
        )
    
    def forward(self, x):
        return self.model(x)
    
    def training_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = nn.functional.cross_entropy(y_hat, y)
        self.log('train_loss', loss)
        return loss
    
    def validation_step(self, batch, batch_idx):
        x, y = batch
        y_hat = self(x)
        loss = nn.functional.cross_entropy(y_hat, y)
        self.log('val_loss', loss)
    
    def configure_optimizers(self):
        return optim.Adam(self.parameters(), lr=self.hparams.learning_rate)

# Training
trainer = pl.Trainer(gpus=1, max_epochs=10)
model = LitModel()
trainer.fit(model, train_loader, val_loader)

### 2. Advanced Lightning Features

# Mixed Precision
trainer = pl.Trainer(precision=16)

# Distributed Training
trainer = pl.Trainer(gpus=2, accelerator='ddp')

# Callbacks
early_stop = pl.callbacks.EarlyStopping(monitor='val_loss')
checkpoint = pl.callbacks.ModelCheckpoint(monitor='val_loss')
trainer = pl.Trainer(callbacks=[early_stop, checkpoint])

# Logging
trainer = pl.Trainer(logger=pl.loggers.TensorBoardLogger('logs/'))

---

## 🔹 Best Practices Summary
1. For GANs: Use spectral norm, progressive growing, and TTUR
2. For VAEs: Monitor both reconstruction and KL divergence terms
3. For RL: Properly normalize rewards and use experience replay
4. For Deployment: Quantize, prune, and export to optimized formats
5. For Maintenance: Use PyTorch Lightning for reproducible experiments

---

### 📌 What's Next?
In Part 6 (Final), we'll cover:
➡️ Advanced Architectures (Graph NNs, Neural ODEs)
➡️ Model Interpretation Techniques
➡️ Production Deployment (TorchServe, Flask API)
➡️ PyTorch Ecosystem (TorchVision, TorchText, TorchAudio)

#PyTorch #DeepLearning #GANs #ReinforcementLearning 🚀

Practice Exercises:
1. Implement WGAN-GP with gradient penalty
2. Train a VAE on MNIST and visualize latent space
3. Build a DQN agent for CartPole environment
4. Quantize a pretrained ResNet and compare accuracy/speed
5. Convert a model to TorchScript and serve with Flask

# WGAN-GP Gradient Penalty
def compute_gradient_penalty(D, real_samples, fake_samples):
    alpha = torch.rand(real_samples.size(0), 1, 1, 1).to(device)
    interpolates = (alpha * real_samples + (1 - alpha) * fake_samples).requires_grad_(True)
    d_interpolates = D(interpolates)
    gradients = torch.autograd.grad(
        outputs=d_interpolates,
        inputs=interpolates,
        grad_outputs=torch.ones_like(d_interpolates),
        create_graph=True,
        retain_graph=True,
        only_inputs=True
    )[0]
    gradients = gradients.view(gradients.size(0), -1)
    gradient_penalty = ((gradients.norm(2, dim=1) - 1) ** 2).mean()
    return gradient_penalty