fast.ai · Parte 2 · Unidad 1 de 5

De Python puro a un framework de entrenamiento propio

Acá arranca la Parte 2: nada de fastai ni vision_learner — se construye, desde la multiplicación de matrices hasta un framework de callbacks completo, todo lo que la Parte 1 usaba como caja negra. Esta unidad condensa cinco lecciones de código denso en tres ideas centrales.

🟢 Por qué esta sección se trata distinto

La Parte 2 (Lecciones 9 a 25, "Deep Learning Foundations to Stable Diffusion") no tiene libro de referencia — son ~30 notebooks de código puro construyendo un mini-framework llamado miniai y, con él, Stable Diffusion desde cero. Es mucho más código que prosa. Estas páginas la condensan en 5 unidades temáticas con los ejemplos más ilustrativos, en vez de una página por lección — el objetivo es cubrir el terreno conceptual completo sin reproducir cada notebook línea por línea.

1. Matmul desde cero, y por qué importa la velocidad

El punto de partida es literal: multiplicar dos matrices con tres loops anidados de Python puro, sobre MNIST.

def matmul(a,b):
    (ar,ac),(br,bc) = a.shape,b.shape
    c = torch.zeros(ar, bc)
    for i in range(ar):
        for j in range(bc):
            for k in range(ac): c[i,j] += a[i,k] * b[k,j]
    return c

Sobre una matriz 5×784 por 784×10, esto tarda ~500ms. El resto de la unidad es una escalera de optimizaciones, cada una reemplazando un loop por una operación vectorizada, midiendo la mejora real en cada paso: Numba (@njit) compila el loop más interno y elimina uno de los tres; broadcasting (a[i,:,None] * b) elimina el segundo, dejando que PyTorch vectorize en C; einsum (torch.einsum('ik,kj->ij', a, b)) expresa la operación completa de forma declarativa; y finalmente, una versión CUDA escrita a mano con @cuda.jit, ejecutando miles de threads en paralelo sobre la GPU.

El resultado medido en el notebook: de ~500ms (Python puro) a ~0.5ms (CUDA) — una mejora de aproximadamente 5 millones de veces sobre todo el dataset. La lección no es "usá broadcasting" per se — es que cada capa de abstracción entre vos y el hardware (Python interpretado → C vectorizado → kernels paralelos en GPU) tiene un costo medible y real, y entender dónde vive ese costo es lo que te permite decidir cuándo vale la pena bajar un nivel.

2. Backprop a mano, capa por capa

Con matmul resuelto, el siguiente paso es implementar el forward y backward pass de una red de dos capas completamente a mano — sin autograd todavía:

def forward_and_backward(inp, targ):
    l1 = lin(inp, w1, b1); l2 = relu(l1); out = lin(l2, w2, b2)
    diff = out[:,0]-targ
    loss = diff.pow(2).mean()
    # backward: la regla de la cadena, escrita explícitamente
    out.g = 2.*diff[:,None] / inp.shape[0]
    lin_grad(l2, out, w2, b2)
    l1.g = (l1>0).float() * l2.g   # derivada del ReLU
    lin_grad(inp, l1, w1, b1)

Este código se verifica comparándolo contra loss.backward() de PyTorch sobre los mismos tensores — y coincide, porque autograd no hace nada distinto: aplica la misma regla de la cadena, automáticamente. El notebook refactoriza esto en tres pasos progresivos: primero clases con backward() propio por capa (Relu, Lin, Mse), después una clase base Module que factoriza el patrón común, y finalmente la versión con nn.Module real de PyTorch — mostrando en código que la "magia" de .backward() es, literalmente, la misma clase Module escrita a mano un párrafo antes.

3. El framework Learner: callbacks en vez de if/else

La pieza más influyente de esta unidad es el diseño del Learner de miniai — básicamente una versión simplificada de por qué Learner/Trainer de fastai y HF están diseñados como están. El problema que resuelve: un training loop necesita soportar mixed precision, logging, early stopping, schedulers de LR, multi-GPU, etc. — si cada feature se agrega como un if dentro del loop, el código se vuelve inmanejable rápido.

La solución es un sistema de callbacks con puntos de enganche nombrados (before_fit, before_batch, after_batch, after_epoch...) más tres excepciones especiales para control de flujo:

class CancelFitException(Exception): pass
class CancelBatchException(Exception): pass
class CancelEpochException(Exception): pass

class SingleBatchCB(Callback):
    order = 1
    def after_batch(self, learn): raise CancelFitException()

Usar excepciones para control de flujo (no para errores) es la parte más elegante del diseño: un callback de "encontrar el learning rate óptimo" puede simplemente raise CancelFitException() cuando el loss explota, y el Learner no necesita saber nada de esa lógica — solo atrapa la excepción en el nivel correspondiente. Con este esqueleto, agregar mixed precision, un progress bar con gráfico en vivo (ProgressCB), tracking de métricas (MetricsCB) o mover tensores a GPU (DeviceCB) es escribir una clase nueva de 5-10 líneas, sin tocar el Learner en absoluto.

🔧 El mismo patrón que ya usás en producción

CancelFitException/CancelBatchException resuelven exactamente el mismo problema que el numeric_gate.py de tu gate anti-alucinación: interrumpir un flujo en curso desde un punto de control externo, sin que la lógica principal necesite saber por qué se interrumpió. Y el diseño de callbacks con order explícito (para controlar en qué secuencia se ejecutan cuando hay varios enganchados al mismo punto) es la misma idea que HuggingFace Trainer usa con sus TrainerCallback — no es casualidad, es el mismo problema de ingeniería resuelto por gente que se leyó el mismo código.

4. Resumen

Cada capa de abstracción entre Python y el hardware (interpretado → vectorizado → CUDA) tiene un costo medible — la escalera matmul lo cuantifica en millones de veces.
Backprop no es magia: es la regla de la cadena aplicada capa por capa, verificable contra autograd de PyTorch porque hace exactamente lo mismo.
Un framework de entrenamiento robusto se construye con callbacks + excepciones de control de flujo, no con ifs acumulados dentro del loop principal.
Este patrón de diseño no es específico de miniai — es el mismo que usan fastai, HF Trainer, y sistemas de control externos como un gate de seguridad.