Modelos generativos: autoencoders, GANs y difusión

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import LinearSegmentedColormap
import torch
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets
from torchvision.utils import make_grid

torch.manual_seed(0)
device = "cuda" if torch.cuda.is_available() else "cpu"
# Del blanco al color "C0": la prenda oscura sobre fondo claro
tinta = LinearSegmentedColormap.from_list("tinta", ["white", "C0"])

entrenamiento = datasets.FashionMNIST(root="datos", train=True,
                                      download=True)
prueba = datasets.FashionMNIST(root="datos", train=False, download=True)
nombres = ["polera", "pantalón", "chompa", "vestido", "abrigo",
           "sandalia", "camisa", "zapatilla", "bolso", "botín"]

g = torch.Generator().manual_seed(0)
idx_ent = torch.randperm(len(entrenamiento), generator=g)[:10000]
idx_val = torch.randperm(len(prueba), generator=g)[:2000]
# Píxeles entre 0 y 1; cada imagen aplanada en un vector de 784
x_ent = entrenamiento.data[idx_ent].float().div(255).view(-1, 784)
x_val = prueba.data[idx_val].float().div(255).view(-1, 784)
y_val = prueba.targets[idx_val]
print("Entrenamiento:", tuple(x_ent.shape), " validación:",
      tuple(x_val.shape))
print(f"Píxeles entre {x_ent.min():.0f} y {x_ent.max():.0f}")

class Autoencoder(nn.Module):
    def __init__(self, latente=16):
        super().__init__()
        self.codificador = nn.Sequential(
            nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, latente))
        self.decodificador = nn.Sequential(
            nn.Linear(latente, 256), nn.ReLU(), nn.Linear(256, 784),
            nn.Sigmoid())

    def forward(self, x):
        return self.decodificador(self.codificador(x))


def con_ruido(x, ruido):
    """Suma ruido gaussiano y deja los píxeles entre 0 y 1."""
    return (x + ruido * torch.randn_like(x)).clamp(0, 1)


def entrenar_ae(modelo, x_ent, x_val, epocas, ruido=0.0, mostrar=True):
    """Entrena para que la salida copie la imagen limpia."""
    modelo.to(device)
    optimizador = torch.optim.Adam(modelo.parameters(), lr=1e-3)
    cargador = DataLoader(TensorDataset(x_ent), batch_size=64,
                          shuffle=True,
                          generator=torch.Generator().manual_seed(0))
    historial = []
    for epoca in range(1, epocas + 1):
        modelo.train()
        for (x,) in cargador:
            x = x.to(device)
            salida = modelo(con_ruido(x, ruido))
            perdida = ((salida - x) ** 2).mean()   # el objetivo es x
            optimizador.zero_grad()
            perdida.backward()
            optimizador.step()
        modelo.eval()
        with torch.no_grad():
            xv = x_val.to(device)
            error = ((modelo(con_ruido(xv, ruido)) - xv) ** 2).mean()
        historial.append(error.item())
        if mostrar:
            print(f"Época {epoca}: error de reconstrucción "
                  f"(validación) {error.item():.4f}")
    return historial


ae = Autoencoder(latente=16)
print("Parámetros:", sum(p.numel() for p in ae.parameters()))
historial = entrenar_ae(ae, x_ent, x_val, epocas=10)

promedio = x_ent.mean(dim=0)
print(f"Error de devolver siempre la prenda promedio: "
      f"{((x_val - promedio) ** 2).mean():.4f}")
print(f"Error del autoencoder: {historial[-1]:.4f}")
print(f"Compresión: 784 píxeles -> 16 números "
      f"({784 // 16} veces menos)")

ae.eval()
with torch.no_grad():
    muestra = x_val[:8]
    copia = ae(muestra.to(device)).cpu()

# Filas alternadas: 4 originales, sus 4 copias, otras 4 y sus copias
filas = [muestra[:4], copia[:4], muestra[4:], copia[4:]]
lote = torch.cat(filas).view(-1, 1, 28, 28)
mosaico = make_grid(lote, nrow=4, padding=2, pad_value=0)
plt.figure(figsize=(4.8, 4.8))
plt.imshow(mosaico[0], cmap=tinta)
plt.yticks([15, 45, 75, 105], ["Original", "Copia", "Original", "Copia"])
plt.xticks([])
plt.grid(False)
plt.title("Copiar a través de 16 números")
plt.show()

def codigo(x):
    with torch.no_grad():
        return ae.codificador(x.to(device))


def decodificar(z):
    with torch.no_grad():
        return ae.decodificador(z).cpu()


def primera(clase):
    """La primera prenda de esa clase en la validación."""
    return x_val[(y_val == clase).nonzero()[0, 0]]


pesos = torch.linspace(0, 1, 7).unsqueeze(1)          # (7, 1)
filas = []
# Pantalón -> vestido y zapatilla -> botín
for a_img, b_img in [(primera(1), primera(3)), (primera(7), primera(9))]:
    z_a, z_b = codigo(a_img), codigo(b_img)
    filas.append((1 - pesos) * a_img + pesos * b_img)
    filas.append(decodificar((1 - pesos) * z_a + pesos * z_b))
print("Códigos del camino:", tuple(((1 - pesos) * z_a).shape))

lote = torch.cat(filas).view(-1, 1, 28, 28)
mosaico = make_grid(lote, nrow=7, padding=2, pad_value=0)
plt.imshow(mosaico[0], cmap=tinta)
plt.yticks([15, 45, 75, 105], ["Píxeles", "Latente", "Píxeles", "Latente"])
plt.xticks([])
plt.grid(False)
plt.title("Dos caminos entre prendas")
plt.show()

def anillo(n):
    """Puntos reales: un anillo de radio 2, un poco borroso."""
    angulo = torch.rand(n) * 2 * torch.pi
    radio = 2 + 0.1 * torch.randn(n)
    return torch.stack([radio * torch.cos(angulo),
                        radio * torch.sin(angulo)], dim=1)


torch.manual_seed(0)
generador = nn.Sequential(nn.Linear(2, 64), nn.ReLU(),
                          nn.Linear(64, 64), nn.ReLU(), nn.Linear(64, 2))
discriminador = nn.Sequential(nn.Linear(2, 64), nn.ReLU(),
                              nn.Linear(64, 64), nn.ReLU(),
                              nn.Linear(64, 1))     # da un logit
# betas=(0.5, ...): Adam recuerda menos los gradientes viejos, lo
# habitual en GANs (no son los β del programa de ruido de más abajo)
opt_g = torch.optim.Adam(generador.parameters(), lr=1e-3,
                         betas=(0.5, 0.999))
opt_d = torch.optim.Adam(discriminador.parameters(), lr=1e-3,
                         betas=(0.5, 0.999))
bce = nn.BCEWithLogitsLoss()          # entropía cruzada binaria
unos, ceros = torch.ones(256, 1), torch.zeros(256, 1)

z_fijo = torch.randn(1000, 2)         # el mismo ruido para cada foto
fotos, perdidas_d, perdidas_g = {}, [], []
for paso in range(5001):
    if paso in (0, 500, 5000):
        with torch.no_grad():
            fotos[paso] = generador(z_fijo)
    if paso == 5000:
        break
    real = anillo(256)
    falso = generador(torch.randn(256, 2))
    # 1) El discriminador: 1 a los reales, 0 a los falsos
    perdida_d = (bce(discriminador(real), unos)
                 + bce(discriminador(falso.detach()), ceros))
    opt_d.zero_grad()
    perdida_d.backward()
    opt_d.step()
    # 2) El generador: que el discriminador diga 1 a sus falsos
    perdida_g = bce(discriminador(falso), unos)
    opt_g.zero_grad()
    perdida_g.backward()
    opt_g.step()
    perdidas_d.append(perdida_d.item())
    perdidas_g.append(perdida_g.item())

with torch.no_grad():
    radios = fotos[5000].norm(dim=1)
    d_real = torch.sigmoid(discriminador(anillo(1000))).mean()
    d_falso = torch.sigmoid(discriminador(fotos[5000])).mean()
print(f"Radio de los puntos generados: {radios.mean():.2f} "
      f"± {radios.std():.2f}")
print(f"D promedio: reales {d_real:.2f}, falsos {d_falso:.2f}")

fig, ejes = plt.subplots(2, 2, figsize=(4.8, 4.8))
reales = anillo(1000)
paneles = [("Datos reales", reales, "C0"),
           ("Paso 0", fotos[0], "C0"),
           ("Paso 500", fotos[500], "C0"),
           ("Paso 5000", fotos[5000], "C1")]
for eje, (titulo, puntos, color) in zip(ejes.flat, paneles):
    if titulo != "Datos reales":
        eje.scatter(reales[:, 0], reales[:, 1], s=2, color="C2",
                    alpha=0.3)
    eje.scatter(puntos[:, 0], puntos[:, 1], s=2, color=color)
    eje.set_title(titulo)
    eje.set_xlim(-3, 3)
    eje.set_ylim(-3, 3)
    eje.set_aspect("equal")
    eje.set_xticks([])
    eje.set_yticks([])
plt.tight_layout()
plt.show()

def suavizar(valores, ventana=100):
    """Promedio móvil: cada punto es la media de 100 pasos."""
    return np.convolve(valores, np.ones(ventana) / ventana, mode="valid")


plt.plot(suavizar(perdidas_d), color="C0", label="Discriminador")
plt.plot(suavizar(perdidas_g), color="C1", label="Generador")
plt.axhline(2 * np.log(2), color="C2", linestyle="--")
plt.axhline(np.log(2), color="C2", linestyle="--")
plt.xlabel("Paso")
plt.ylabel("Pérdida")
plt.title("Las pérdidas de la GAN del anillo")
plt.legend()
plt.show()

# La primera vez descarga unos 11 MB; después queda en «datos»
mnist = datasets.MNIST(root="datos", train=True, download=True)
# Píxeles de -1 a 1, aplanados
x_mnist = mnist.data.float().div(255).mul(2).sub(1).view(-1, 784)

torch.manual_seed(0)
ruido_dim = 64
gen_dig = nn.Sequential(
    nn.Linear(ruido_dim, 256), nn.LeakyReLU(0.2),
    nn.Linear(256, 512), nn.LeakyReLU(0.2),
    nn.Linear(512, 784), nn.Tanh()).to(device)
dis_dig = nn.Sequential(
    nn.Linear(784, 512), nn.LeakyReLU(0.2),
    nn.Linear(512, 256), nn.LeakyReLU(0.2),
    nn.Linear(256, 1)).to(device)
opt_g = torch.optim.Adam(gen_dig.parameters(), lr=2e-4, betas=(0.5, 0.999))
opt_d = torch.optim.Adam(dis_dig.parameters(), lr=2e-4, betas=(0.5, 0.999))
cargador = DataLoader(TensorDataset(x_mnist), batch_size=128, shuffle=True,
                      drop_last=True,
                      generator=torch.Generator().manual_seed(0))

for epoca in range(1, 21):
    suma_d = suma_g = 0.0
    for (real,) in cargador:
        real = real.to(device)
        n = len(real)
        unos = torch.ones(n, 1, device=device)
        ceros = torch.zeros(n, 1, device=device)
        falso = gen_dig(torch.randn(n, ruido_dim, device=device))
        perdida_d = (bce(dis_dig(real), unos)
                     + bce(dis_dig(falso.detach()), ceros))
        opt_d.zero_grad()
        perdida_d.backward()
        opt_d.step()
        perdida_g = bce(dis_dig(falso), unos)
        opt_g.zero_grad()
        perdida_g.backward()
        opt_g.step()
        suma_d += perdida_d.item()
        suma_g += perdida_g.item()
    if epoca % 5 == 0:
        print(f"Época {epoca}: pérdida D {suma_d / len(cargador):.3f}, "
              f"G {suma_g / len(cargador):.3f}")

with torch.no_grad():
    digitos = gen_dig(torch.randn(36, ruido_dim, device=device)).cpu()
digitos = digitos.add(1).div(2).view(-1, 1, 28, 28)   # de vuelta a 0..1
mosaico = make_grid(digitos, nrow=6, padding=2, pad_value=0)
plt.figure(figsize=(4.8, 4.8))
plt.imshow(mosaico[0], cmap=tinta)
plt.axis("off")
plt.title("36 dígitos que no escribió nadie")
plt.show()

rng = np.random.default_rng(42)


def espiral(n):
    """Puntos sobre una espiral de dos vueltas, con un poco de ruido."""
    angulo = rng.uniform(np.pi / 2, 4 * np.pi, n)
    radio = angulo / (2 * np.pi)
    xy = np.stack([radio * np.cos(angulo), radio * np.sin(angulo)], axis=1)
    xy += rng.normal(0, 0.04, size=(n, 2))
    return torch.tensor(xy, dtype=torch.float32)


datos_2d = espiral(2000)

T = 100
# beta[0] = 0: el paso 0 son los datos sin ruido
beta = torch.cat([torch.zeros(1), torch.linspace(1e-4, 0.1, T)])
alfa = 1 - beta
alfa_barra = torch.cumprod(alfa, dim=0)


def agregar_ruido(x0, t, eps):
    """La ida en un solo golpe: x_t a partir de x_0 y el ruido eps."""
    ab = alfa_barra[t].view(-1, 1)
    return ab.sqrt() * x0 + (1 - ab).sqrt() * eps


for t in [0, 10, 20, 50, 100]:
    print(f"t = {t:3d}: señal {alfa_barra[t].sqrt():.2f}, "
          f"ruido {(1 - alfa_barra[t]).sqrt():.2f}")

eps = torch.randn_like(datos_2d)
fig, ejes = plt.subplots(2, 2, figsize=(4.8, 4.8))
for eje, t in zip(ejes.flat, [0, 10, 20, 100]):
    xt = agregar_ruido(datos_2d, torch.full((2000,), t), eps)
    eje.scatter(xt[:, 0], xt[:, 1], s=1, color="C0")
    eje.set_title(f"t = {t}")
    eje.set_xlim(-3.2, 3.2)
    eje.set_ylim(-3.2, 3.2)
    eje.set_aspect("equal")
    eje.set_xticks([])
    eje.set_yticks([])
plt.tight_layout()
plt.show()

def incrustar(t):
    """Codifica el paso t con senos y cosenos de 8 frecuencias."""
    frecuencias = 2.0 ** torch.arange(8)
    a = t.float().unsqueeze(1) / T * torch.pi * frecuencias
    return torch.cat([a.sin(), a.cos()], dim=1)


torch.manual_seed(0)
# SiLU: una ReLU suave, x · sigmoide(x)
red = nn.Sequential(nn.Linear(2 + 16, 128), nn.SiLU(),
                    nn.Linear(128, 128), nn.SiLU(),
                    nn.Linear(128, 128), nn.SiLU(),
                    nn.Linear(128, 2))         # el ruido estimado
pasos = 15000
optimizador = torch.optim.Adam(red.parameters(), lr=2e-3)
# Programador de tasa (tema 19): la tasa baja suavemente hasta 0
programador = torch.optim.lr_scheduler.CosineAnnealingLR(optimizador,
                                                         pasos)
perdidas = []
for paso in range(pasos):
    x0 = datos_2d[torch.randint(0, 2000, (512,))]
    t = torch.randint(1, T + 1, (512,))
    eps = torch.randn_like(x0)
    xt = agregar_ruido(x0, t, eps)
    eps_estimado = red(torch.cat([xt, incrustar(t)], dim=1))
    perdida = ((eps_estimado - eps) ** 2).mean()
    optimizador.zero_grad()
    perdida.backward()
    optimizador.step()
    programador.step()
    perdidas.append(perdida.item())

print(f"Pérdida media, primeros 500 pasos: {np.mean(perdidas[:500]):.3f}")
print(f"Pérdida media, últimos 500 pasos:  {np.mean(perdidas[-500:]):.3f}")

@torch.no_grad()
def muestrear(n, guardar=(100, 20, 10, 0)):
    x = torch.randn(n, 2)                 # ruido puro: el paso T
    fotos = {T: x.clone()}
    for t in range(T, 0, -1):
        tt = torch.full((n,), t)
        e = red(torch.cat([x, incrustar(tt)], dim=1))
        x = (x - beta[t] / (1 - alfa_barra[t]).sqrt() * e) / alfa[t].sqrt()
        if t > 1:
            x = x + beta[t].sqrt() * torch.randn_like(x)
        if t - 1 in guardar:
            fotos[t - 1] = x.clone()
    return fotos


torch.manual_seed(1)
vuelta = muestrear(2000)
fig, ejes = plt.subplots(2, 2, figsize=(4.8, 4.8))
for eje, t in zip(ejes.flat, [100, 20, 10, 0]):
    color = "C1" if t == 0 else "C0"
    eje.scatter(vuelta[t][:, 0], vuelta[t][:, 1], s=1, color=color)
    eje.set_title(f"t = {t}")
    eje.set_xlim(-3.2, 3.2)
    eje.set_ylim(-3.2, 3.2)
    eje.set_aspect("equal")
    eje.set_xticks([])
    eje.set_yticks([])
plt.tight_layout()
plt.show()

from diffusers import DDPMPipeline, DDIMPipeline, DDIMScheduler

# La primera vez descarga unos 143 MB
ddpm = DDPMPipeline.from_pretrained("google/ddpm-cifar10-32")
ddpm.set_progress_bar_config(disable=True)
ddim = DDIMPipeline(unet=ddpm.unet,
                    scheduler=DDIMScheduler.from_config(
                        ddpm.scheduler.config))
ddim.set_progress_bar_config(disable=True)
print("Parámetros de la U-Net:",
      f"{sum(p.numel() for p in ddpm.unet.parameters()):,}")
print("Pasos con que se entrenó:", ddpm.scheduler.config.num_train_timesteps)

g = torch.Generator().manual_seed(0)
imgs_ddpm = ddpm(batch_size=4, num_inference_steps=50, generator=g,
                 output_type="np").images            # (4, 32, 32, 3)
g = torch.Generator().manual_seed(0)
imgs_ddim = ddim(batch_size=4, num_inference_steps=50, generator=g,
                 output_type="np").images
print("Forma:", imgs_ddpm.shape, " valores entre 0 y 1")

# Cuatro de DDPM arriba y cuatro de DDIM abajo
lote = torch.tensor(np.concatenate([imgs_ddpm, imgs_ddim]))
mosaico = make_grid(lote.permute(0, 3, 1, 2), nrow=4, padding=2,
                    pad_value=1)
plt.imshow(mosaico.permute(1, 2, 0))
plt.yticks([17, 51], ["DDPM", "DDIM"])
plt.xticks([])
plt.grid(False)
plt.title("50 pasos con la misma red")

plt.show()