Entrenando una GAN con PyTorch

# CPU (funciona en cualquier máquina)
pip install torch torchvision matplotlib

# GPU CUDA 12.1
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import make_grid
import matplotlib.pyplot as plt
import numpy as np
import os
import time

# Reproducibilidad
SEED = 42
torch.manual_seed(SEED)
np.random.seed(SEED)

# Device
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"PyTorch {torch.__version__} | Device: {device}")

# ── Hiperparámetros ──────────────────────────────────────
LATENT_DIM = 100       # Dimensión del vector de ruido z
IMG_SIZE = 28          # MNIST: 28x28
IMG_CHANNELS = 1       # Escala de grises
IMG_PIXELS = IMG_SIZE * IMG_SIZE * IMG_CHANNELS  # 784

BATCH_SIZE = 128
EPOCHS = 50
LR_G = 2e-4            # Learning rate del Generador
LR_D = 2e-4            # Learning rate del Discriminador
BETA1 = 0.5            # Adam beta1 (estándar en GANs)
BETA2 = 0.999          # Adam beta2

# Directorio para guardar resultados
os.makedirs("gan_results", exist_ok=True)

# ── Transformaciones ─────────────────────────────────────
# ToTensor() convierte [0, 255] → [0, 1]
# Normalize((0.5,), (0.5,)) mapea [0, 1] → [-1, 1]
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])

# ── Descargar y cargar MNIST ─────────────────────────────
train_dataset = datasets.MNIST(
    root='./data',
    train=True,
    transform=transform,
    download=True
)

train_loader = DataLoader(
    train_dataset,
    batch_size=BATCH_SIZE,
    shuffle=True,
    drop_last=True,   # Descartar último batch incompleto
    num_workers=2,
    pin_memory=True if device.type == 'cuda' else False,
)

print(f"Dataset: {len(train_dataset):,} imágenes")
print(f"Batches por epoch: {len(train_loader)}")
print(f"Imagen shape: {train_dataset[0][0].shape}")
print(f"Rango: [{train_dataset[0][0].min():.1f}, {train_dataset[0][0].max():.1f}]")

def show_images(images, nrow=8, title=""):
    """Muestra un grid de imágenes (tensor normalizado a [-1,1])."""
    # Desnormalizar: [-1, 1] → [0, 1]
    images = (images + 1) / 2
    images = images.clamp(0, 1)
    grid = make_grid(images, nrow=nrow, padding=2)
    plt.figure(figsize=(10, 10))
    plt.imshow(grid.permute(1, 2, 0).cpu().numpy(), cmap='gray')
    plt.title(title, fontsize=14)
    plt.axis('off')
    plt.tight_layout()
    plt.show()

# Visualizar un batch de imágenes reales
real_batch = next(iter(train_loader))[0][:64]
show_images(real_batch, title="MNIST — Imágenes reales")

class Generator(nn.Module):
    """
    Generador MLP: z (100,) → imagen (1, 28, 28).
    Arquitectura: Linear → LeakyReLU → BN, repetido, → Tanh.
    """
    def __init__(self, latent_dim=LATENT_DIM, img_pixels=IMG_PIXELS):
        super().__init__()
        self.net = nn.Sequential(
            # Bloque 1: 100 → 256
            nn.Linear(latent_dim, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.BatchNorm1d(256),

            # Bloque 2: 256 → 512
            nn.Linear(256, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.BatchNorm1d(512),

            # Bloque 3: 512 → 1024
            nn.Linear(512, 1024),
            nn.LeakyReLU(0.2, inplace=True),
            nn.BatchNorm1d(1024),

            # Capa de salida: 1024 → 784, activación Tanh
            nn.Linear(1024, img_pixels),
            nn.Tanh(),  # Output en [-1, 1]
        )

    def forward(self, z):
        """z: (batch, latent_dim) → img: (batch, 1, 28, 28)"""
        flat = self.net(z)                         # (B, 784)
        img = flat.view(-1, IMG_CHANNELS, IMG_SIZE, IMG_SIZE)  # (B, 1, 28, 28)
        return img

# Crear generador
G = Generator().to(device)

# Verificar shape
z_test = torch.randn(4, LATENT_DIM, device=device)
fake_test = G(z_test)
print(f"Generador - Input: {z_test.shape} → Output: {fake_test.shape}")
print(f"Parámetros G: {sum(p.numel() for p in G.parameters()):,}")

class Discriminator(nn.Module):
    """
    Discriminador MLP: imagen (1, 28, 28) → probabilidad [0, 1].
    Arquitectura: Flatten → Linear → LeakyReLU → Dropout, repetido, → Sigmoid.
    """
    def __init__(self, img_pixels=IMG_PIXELS):
        super().__init__()
        self.net = nn.Sequential(
            # Flatten: (B, 1, 28, 28) → (B, 784)
            nn.Flatten(),

            # Bloque 1: 784 → 512
            nn.Linear(img_pixels, 512),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Dropout(0.3),

            # Bloque 2: 512 → 256
            nn.Linear(512, 256),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Dropout(0.3),

            # Bloque 3: 256 → 128
            nn.Linear(256, 128),
            nn.LeakyReLU(0.2, inplace=True),
            nn.Dropout(0.3),

            # Salida: 128 → 1 (probabilidad)
            nn.Linear(128, 1),
            nn.Sigmoid(),
        )

    def forward(self, img):
        """img: (batch, 1, 28, 28) → p: (batch, 1)"""
        return self.net(img)

# Crear discriminador
D = Discriminator().to(device)

# Verificar shape
pred_test = D(fake_test)
print(f"Discriminador - Input: {fake_test.shape} → Output: {pred_test.shape}")
print(f"Parámetros D: {sum(p.numel() for p in D.parameters()):,}")

def weights_init(m):
    """Inicialización de pesos según DCGAN paper."""
    classname = m.__class__.__name__
    if classname.find('Linear') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
        if m.bias is not None:
            nn.init.constant_(m.bias.data, 0)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)

# Aplicar a G y D
G.apply(weights_init)
D.apply(weights_init)
print("✓ Pesos inicializados (N(0, 0.02))")

# ── Loss: Binary Cross-Entropy ───────────────────────────
criterion = nn.BCELoss()

# ── Optimizadores separados para G y D ───────────────────
optimizer_G = optim.Adam(G.parameters(), lr=LR_G, betas=(BETA1, BETA2))
optimizer_D = optim.Adam(D.parameters(), lr=LR_D, betas=(BETA1, BETA2))

# ── Labels fijos ─────────────────────────────────────────
real_label = 1.0
fake_label = 0.0

# ── Vector z fijo para visualización ─────────────────────
# Lo usaremos para ver cómo evolucionan las generaciones
fixed_noise = torch.randn(64, LATENT_DIM, device=device)

# ── Logging ──────────────────────────────────────────────
G_losses, D_losses = [], []
D_real_acc, D_fake_acc = [], []

print("Iniciando entrenamiento...")
start_time = time.time()

for epoch in range(EPOCHS):
    g_loss_epoch, d_loss_epoch = 0, 0
    d_real_epoch, d_fake_epoch = 0, 0

    for i, (real_imgs, _) in enumerate(train_loader):
        batch_size = real_imgs.size(0)
        real_imgs = real_imgs.to(device)

        # ═══════════════════════════════════════════════════
        # FASE 1: Entrenar el Discriminador
        # Objetivo: maximizar log(D(x)) + log(1 - D(G(z)))
        # ═══════════════════════════════════════════════════
        D.zero_grad()

        # ── 1a. Imágenes REALES → D debe decir "real" (1) ──
        labels_real = torch.full((batch_size, 1), real_label,
                                  device=device)
        output_real = D(real_imgs)
        loss_D_real = criterion(output_real, labels_real)

        # ── 1b. Imágenes FAKE → D debe decir "fake" (0) ────
        z = torch.randn(batch_size, LATENT_DIM, device=device)
        fake_imgs = G(z)
        labels_fake = torch.full((batch_size, 1), fake_label,
                                  device=device)
        output_fake = D(fake_imgs.detach())  # ¡DETACH! No queremos gradientes en G
        loss_D_fake = criterion(output_fake, labels_fake)

        # ── 1c. Loss total de D y backward ──────────────────
        loss_D = (loss_D_real + loss_D_fake) / 2
        loss_D.backward()
        optimizer_D.step()

        # ═══════════════════════════════════════════════════
        # FASE 2: Entrenar el Generador
        # Objetivo: maximizar log(D(G(z))) → "engañar a D"
        # ═══════════════════════════════════════════════════
        G.zero_grad()

        # ── 2a. Generar fakes y pedir a D que diga "real" ──
        z = torch.randn(batch_size, LATENT_DIM, device=device)
        fake_imgs = G(z)
        labels_real_for_G = torch.full((batch_size, 1), real_label,
                                        device=device)
        output_G = D(fake_imgs)  # SIN detach: queremos gradientes en G
        loss_G = criterion(output_G, labels_real_for_G)

        loss_G.backward()
        optimizer_G.step()

        # ── Acumular métricas ────────────────────────────────
        g_loss_epoch += loss_G.item()
        d_loss_epoch += loss_D.item()
        d_real_epoch += output_real.mean().item()
        d_fake_epoch += output_fake.mean().item()

    # ── Promedios del epoch ──────────────────────────────────
    n = len(train_loader)
    G_losses.append(g_loss_epoch / n)
    D_losses.append(d_loss_epoch / n)
    D_real_acc.append(d_real_epoch / n)
    D_fake_acc.append(d_fake_epoch / n)

    elapsed = time.time() - start_time
    if (epoch + 1) % 5 == 0 or epoch == 0:
        print(f"Epoch [{epoch+1:3d}/{EPOCHS}] | "
              f"D_loss: {D_losses[-1]:.4f} | G_loss: {G_losses[-1]:.4f} | "
              f"D(x): {D_real_acc[-1]:.3f} | D(G(z)): {D_fake_acc[-1]:.3f} | "
              f"Time: {elapsed:.0f}s")

    # ── Guardar grid de imágenes cada 10 epochs ─────────────
    if (epoch + 1) % 10 == 0 or epoch == 0:
        with torch.no_grad():
            sample = G(fixed_noise)
            sample = (sample + 1) / 2  # [-1,1] → [0,1]
            grid = make_grid(sample, nrow=8, padding=2)
            plt.figure(figsize=(8, 8))
            plt.imshow(grid.permute(1, 2, 0).cpu().numpy(), cmap='gray')
            plt.title(f'Epoch {epoch+1}')
            plt.axis('off')
            plt.savefig(f'gan_results/epoch_{epoch+1:03d}.png',
                       bbox_inches='tight', dpi=100)
            plt.close()

print(f"\n✓ Entrenamiento completo en {time.time()-start_time:.0f}s")

fig, axes = plt.subplots(1, 2, figsize=(14, 5))

# ── Panel 1: Losses ──
axes[0].plot(G_losses, label='Generator', color='#6c5ce7', linewidth=2)
axes[0].plot(D_losses, label='Discriminator', color='#fd79a8', linewidth=2)
axes[0].axhline(y=np.log(2), color='rgba(253,203,110,.5)', 
                linestyle='--', label=f'Equilibrio (ln2≈{np.log(2):.3f})')
axes[0].set_xlabel('Epoch')
axes[0].set_ylabel('Loss')
axes[0].set_title('GAN Losses')
axes[0].legend()
axes[0].grid(alpha=0.15)

# ── Panel 2: Confianza de D ──
axes[1].plot(D_real_acc, label='D(x) — reales', color='#00b894', linewidth=2)
axes[1].plot(D_fake_acc, label='D(G(z)) — fakes', color='#e17055', linewidth=2)
axes[1].axhline(y=0.5, color='rgba(255,255,255,.2)', linestyle='--', label='Equilibrio (0.5)')
axes[1].set_xlabel('Epoch')
axes[1].set_ylabel('Probabilidad media')
axes[1].set_title('Confianza del Discriminador')
axes[1].legend()
axes[1].grid(alpha=0.15)
axes[1].set_ylim(0, 1)

plt.tight_layout()
plt.savefig('gan_results/training_curves.png', dpi=150, bbox_inches='tight')
plt.show()

def show_evolution(G, fixed_noise, epochs_to_show=[1, 10, 25, 50]):
    """Genera imágenes con el mismo ruido guardado en diferentes checkpoints."""
    # Esta función requiere haber guardado modelos en cada epoch
    # Aquí mostramos cómo generar un grid al final del entrenamiento
    G.eval()
    with torch.no_grad():
        fake = G(fixed_noise)
        fake = (fake + 1) / 2  # [-1,1] → [0,1]
    G.train()

    grid = make_grid(fake, nrow=8, padding=2)
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(grid.permute(1, 2, 0).cpu().numpy(), cmap='gray')
    ax.set_title('Imágenes generadas (epoch final)', fontsize=14)
    ax.axis('off')
    plt.tight_layout()
    plt.savefig('gan_results/final_generation.png', dpi=150, bbox_inches='tight')
    plt.show()

show_evolution(G, fixed_noise)

def interpolate_latent(G, z1, z2, n_steps=10):
    """Interpola linealmente entre z1 y z2 en el espacio latente."""
    G.eval()
    alphas = torch.linspace(0, 1, n_steps, device=device)
    interpolations = []

    with torch.no_grad():
        for alpha in alphas:
            z = (1 - alpha) * z1 + alpha * z2
            img = G(z.unsqueeze(0))
            interpolations.append(img)

    images = torch.cat(interpolations, dim=0)
    images = (images + 1) / 2
    G.train()
    return images

# Dos puntos aleatorios del espacio latente
z1 = torch.randn(LATENT_DIM, device=device)
z2 = torch.randn(LATENT_DIM, device=device)

interp = interpolate_latent(G, z1, z2, n_steps=12)

grid = make_grid(interp, nrow=12, padding=2)
plt.figure(figsize=(15, 2))
plt.imshow(grid.permute(1, 2, 0).cpu().numpy(), cmap='gray')
plt.title('Interpolación en espacio latente: z₁ → z₂', fontsize=12)
plt.axis('off')
plt.tight_layout()
plt.savefig('gan_results/interpolation.png', dpi=150, bbox_inches='tight')
plt.show()

class DCGenerator(nn.Module):
    """
    DCGAN Generator: z (100,) → imagen (1, 28, 28).
    Usa ConvTranspose2d para upsampling progresivo.
    """
    def __init__(self, latent_dim=LATENT_DIM, feature_maps=128):
        super().__init__()
        self.net = nn.Sequential(
            # z: (B, 100) → (B, 256, 7, 7)
            nn.ConvTranspose2d(latent_dim, feature_maps * 2, 7, 1, 0, bias=False),
            nn.BatchNorm2d(feature_maps * 2),
            nn.ReLU(True),

            # (B, 256, 7, 7) → (B, 128, 14, 14)
            nn.ConvTranspose2d(feature_maps * 2, feature_maps, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_maps),
            nn.ReLU(True),

            # (B, 128, 14, 14) → (B, 1, 28, 28)
            nn.ConvTranspose2d(feature_maps, IMG_CHANNELS, 4, 2, 1, bias=False),
            nn.Tanh(),
        )

    def forward(self, z):
        # z: (B, 100) → (B, 100, 1, 1) para ConvTranspose2d
        z = z.view(-1, LATENT_DIM, 1, 1)
        return self.net(z)

# Crear y verificar
G_dc = DCGenerator().to(device)
G_dc.apply(weights_init)
z_test = torch.randn(4, LATENT_DIM, device=device)
out = G_dc(z_test)
print(f"DCGAN Generator: {z_test.shape} → {out.shape}")
print(f"Parámetros: {sum(p.numel() for p in G_dc.parameters()):,}")

class DCDiscriminator(nn.Module):
    """
    DCGAN Discriminator: imagen (1, 28, 28) → probabilidad [0, 1].
    Usa Conv2d con stride para downsampling progresivo.
    """
    def __init__(self, feature_maps=128):
        super().__init__()
        self.net = nn.Sequential(
            # (B, 1, 28, 28) → (B, 128, 14, 14) — sin BN en la primera capa
            nn.Conv2d(IMG_CHANNELS, feature_maps, 4, 2, 1, bias=False),
            nn.LeakyReLU(0.2, inplace=True),

            # (B, 128, 14, 14) → (B, 256, 7, 7)
            nn.Conv2d(feature_maps, feature_maps * 2, 4, 2, 1, bias=False),
            nn.BatchNorm2d(feature_maps * 2),
            nn.LeakyReLU(0.2, inplace=True),

            # (B, 256, 7, 7) → (B, 1, 1, 1)
            nn.Conv2d(feature_maps * 2, 1, 7, 1, 0, bias=False),
            nn.Sigmoid(),
        )

    def forward(self, img):
        return self.net(img).view(-1, 1)  # (B, 1)

# Crear y verificar
D_dc = DCDiscriminator().to(device)
D_dc.apply(weights_init)
pred = D_dc(out.detach())
print(f"DCGAN Discriminator: {out.shape} → {pred.shape}")
print(f"Parámetros: {sum(p.numel() for p in D_dc.parameters()):,}")

# ── Usar los modelos DCGAN ───────────────────────────────
G = G_dc
D = D_dc

optimizer_G = optim.Adam(G.parameters(), lr=LR_G, betas=(BETA1, BETA2))
optimizer_D = optim.Adam(D.parameters(), lr=LR_D, betas=(BETA1, BETA2))

# ── Reutilizar el mismo training loop del paso 6 ────────
# (copia el loop anterior, o refactorízalo en una función)
# El resultado con DCGAN típicamente:
# - Converge más rápido (~20 epochs vs ~50)
# - Produce dígitos más nítidos
# - Captura mejor los detalles espaciales (trazos, curvas)

# ── Con label smoothing ──────────────────────────────────
# En vez de:
#   labels_real = torch.ones(batch_size, 1)    → 1.0
#   labels_fake = torch.zeros(batch_size, 1)   → 0.0

# Usamos:
labels_real = torch.FloatTensor(batch_size, 1).uniform_(0.9, 1.0).to(device)
labels_fake = torch.FloatTensor(batch_size, 1).uniform_(0.0, 0.1).to(device)

# Esto hace que D no pueda ser "100% seguro", lo que:
# 1. Previene que D domine a G
# 2. Regulariza D de forma implícita
# 3. Reduce el riesgo de gradientes saturados