Autograd y el loop de optimización, con nombre y apellido

1. El DAG dinámico: por qué PyTorch permite control de flujo

torch.autograd mantiene, mientras se ejecuta el forward pass, un grafo acíclico dirigido (DAG) de objetos Function — las hojas son los tensores de entrada, la raíz es el tensor de salida (típicamente el loss). Cada operación queda registrada como un nodo con su función de derivada correspondiente, accesible vía .grad_fn:

z = torch.matmul(x, w) + b
loss = F.binary_cross_entropy_with_logits(z, y)

print(z.grad_fn)     # <AddBackward0 object>
print(loss.grad_fn)  # <BinaryCrossEntropyWithLogitsBackward0 object>

El detalle que la documentación oficial marca como clave: este grafo se reconstruye desde cero en cada forward pass. No hay un grafo estático definido de antemano (como en TensorFlow 1.x o en Theano) — cada llamada a .backward() consume el grafo actual, y el próximo forward pass arma uno nuevo. Esto es exactamente lo que permite usar if, loops y control de flujo arbitrario de Python dentro de un modelo: la forma, el tamaño y las operaciones pueden cambiar en cada iteración sin ningún problema, porque no hay un grafo fijo que reconciliar.

2. requires_grad, no_grad y detach

Solo los tensores marcados con requires_grad=True acumulan historial de gradiente — y solo se puede leer .grad en los leaf nodes del grafo (los tensores de entrada, no los resultados intermedios). Para inferencia, donde no hace falta ningún gradiente, dos formas equivalentes de desactivarlo:

with torch.no_grad():
    z = torch.matmul(x, w) + b   # z.requires_grad == False acá dentro

z_det = z.detach()                # mismo efecto, sobre un tensor puntual

Dos razones prácticas para desactivarlo: marcar partes del modelo como parámetros congelados (exactamente lo que pasa con el modelo base cuando entrenás solo un adapter LoRA), y acelerar el forward pass cuando no hace falta backward — construir y mantener el grafo tiene un costo real de memoria y cómputo que no vale la pena pagar en pura inferencia.

3. Los gradientes se acumulan — por diseño, no por bug

Cada llamada a .backward() suma el gradiente calculado al valor que ya había en .grad, en vez de reemplazarlo. Es intencional (útil para acumular gradientes de varios mini-batches antes de un solo step, una técnica real para simular batches más grandes de los que entran en memoria), pero significa que sin limpiar explícitamente entre steps, los gradientes de iteraciones distintas se mezclan:

optimizer.zero_grad()    # 1. limpiar gradientes de la iteración anterior
loss.backward()          # 2. calcular gradientes nuevos, sumándolos a .grad
optimizer.step()         # 3. ajustar parámetros usando esos gradientes

Estos tres pasos, en este orden exacto (o con zero_grad() al final, algunas variantes lo ubican ahí), son el corazón de cualquier training loop de PyTorch que vayas a ver — desde un nn.Linear de dos capas hasta el fine-tuning de un LLM de 70B parámetros con LoRA.

4. El loop de optimización completo

Con loss, optimizer y los tres pasos de arriba, el loop completo de la documentación oficial:

def train_loop(dataloader, model, loss_fn, optimizer):
    model.train()
    for batch, (X, y) in enumerate(dataloader):
        pred = model(X)
        loss = loss_fn(pred, y)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

def test_loop(dataloader, model, loss_fn):
    model.eval()
    test_loss, correct = 0, 0
    with torch.no_grad():
        for X, y in dataloader:
            pred = model(X)
            test_loss += loss_fn(pred, y).item()
            correct += (pred.argmax(1) == y).type(torch.float).sum().item()

nn.CrossEntropyLoss combina internamente LogSoftmax + NLLLoss (Negative Log Likelihood) — por eso el modelo devuelve logits crudos, sin softmax aplicado a mano, y es la loss la que se encarga de normalizarlos. Entre los optimizadores disponibles (torch.optim), SGD es el más simple; Adam y RMSProp suelen converger más rápido en la práctica para la mayoría de arquitecturas modernas.

5. model.train() vs model.eval(): un interruptor, no cosmética

La documentación oficial insiste en un detalle que es fácil pasar por alto: llamar model.train() antes de entrenar y model.eval() antes de evaluar no es una convención cosmética. Capas como BatchNorm y Dropout se comportan de forma distinta en cada modo — BatchNorm usa estadísticas del batch actual en training pero un promedio acumulado en eval; Dropout apaga neuronas al azar en training pero las deja todas activas en eval. Olvidarse de este switch produce resultados de inferencia inconsistentes y difíciles de diagnosticar, porque el modelo "funciona", solo que con números ligeramente distintos cada vez de una forma que no tiene una causa obvia hasta que se revisa este detalle puntual.

6. Resumen

1. Autograd reconstruye el grafo de cómputo en cada forward pass — dinámico, no estático, lo que permite control de flujo arbitrario de Python dentro del modelo.
2. no_grad()/detach() desactivan el tracking de gradientes — para congelar parámetros o acelerar inferencia pura.
3. Los gradientes se acumulan por diseño; zero_grad() es obligatorio en cada iteración salvo que quieras acumular a propósito.
4. El loop de optimización son tres pasos fijos: backward(), step(), zero_grad() — el mismo patrón en cualquier escala de modelo.
5. model.train()/model.eval() cambian el comportamiento real de BatchNorm y Dropout — no es un detalle estético, es funcional.