Fine-tuning

El problema

Tienes dos tareas de clasificación con las mismas entradas: 50 lecturas numéricas por muestra. Los datos son sintéticos. Las lecturas dependen de cuatro factores ocultos y de diez fuentes de interferencia. La clase de cada tarea es una combinación distinta de los mismos cuatro factores.

Para entrenar una red 50 → 32 → 32 → 4 en la tarea destino solo hay 20 ejemplos. La tarea fuente, en cambio, tiene datos de sobra para entrenar una red con la misma arquitectura.

Cuerpo y cabeza

La función generar, en la celda plegada, produce las dos tareas a la vez. Cada muestra tiene cuatro factores ocultos y diez valores de interferencia. Las 50 lecturas mezclan ambos y pasan por tanh. La clase fuente dice qué factor tiene el mayor valor absoluto y con qué signo. La clase destino sale de otra combinación de los mismos factores, con pesos aleatorios.

ejemplos_destino genera muestras nuevas y toma las primeras de cada clase destino. Todas las semillas son distintas, así que entrenamiento y prueba nunca comparten muestras.

import numpy as np
import torch

rng = np.random.default_rng(0)
pesos_fuente = np.hstack([np.eye(4), -np.eye(4)])  # columnas 0-3 con signo +, 4-7 con signo −
mezcla_factores = rng.normal(size=(4, 50)) / 2
mezcla_interferencia = rng.normal(size=(10, 50)) / np.sqrt(10)
pesos_destino = rng.normal(size=(4, 4))


def generar(n, semilla, pesos_clase_fuente=pesos_fuente):
    g = np.random.default_rng(semilla)
    factores = g.normal(size=(n, 4))
    interferencia = g.normal(scale=2.0, size=(n, 10))
    X = np.tanh(factores @ mezcla_factores + interferencia @ mezcla_interferencia)
    X = X + g.normal(scale=0.1, size=X.shape)
    y_fuente = (factores @ pesos_clase_fuente).argmax(axis=1)  # columna con el mayor valor
    y_destino = (factores @ pesos_destino).argmax(axis=1)
    return (
        torch.tensor(X, dtype=torch.float32),
        torch.tensor(y_fuente),
        torch.tensor(y_destino),
    )


def ejemplos_destino(por_clase, semilla):
    X, _, y = generar(4000, semilla)
    indices = torch.cat(  # posiciones de las primeras por_clase muestras de cada clase
        [torch.nonzero(y == clase).flatten()[:por_clase] for clase in range(4)]
    )
    return X[indices], y[indices]

La segunda celda plegada genera las 20 000 muestras fuente (X_fuente, y_fuente), las 2000 de prueba con las etiquetas de ambas tareas (X_prueba, y_fuente_prueba, y_destino_prueba) y los 20 ejemplos destino (X_pocos, y_pocos).

X_fuente, y_fuente, _ = generar(20000, semilla=1)
X_prueba, y_fuente_prueba, y_destino_prueba = generar(2000, semilla=2)
X_pocos, y_pocos = ejemplos_destino(5, semilla=10)

Pregunta. ¿Las ocho clases fuente tienen un número parecido de ejemplos? ¿Y las cuatro clases destino en la prueba?

Una red de clasificación se puede dividir en dos partes. El cuerpo son todas las capas menos la última: transforma las 50 lecturas en 32 características. La cabeza es la última capa lineal: convierte esas características en un logit por clase. El número de clases solo cambia la forma de la cabeza. La clase Clasificador guarda cada parte en su propio atributo, y nn.Sequential encadena las capas del cuerpo en orden. crear_modelo fija la semilla, así que cada llamada con el mismo número de clases devuelve los mismos pesos iniciales.

from torch import nn


class Clasificador(nn.Module):
    def __init__(self, n_clases):
        super().__init__()
        self.cuerpo = nn.Sequential(
            nn.Linear(50, 32), nn.ReLU(), nn.Linear(32, 32), nn.ReLU()
        )
        self.cabeza = nn.Linear(32, n_clases)

    def forward(self, x):
        return self.cabeza(self.cuerpo(x))


def crear_modelo(n_clases):
    torch.manual_seed(0)
    return Clasificador(n_clases)

entrenar usa todos los ejemplos en cada paso, sin DataLoader. Con 20 ejemplos no hace falta dividirlos en batches, y con 20 000 ejemplos de 50 lecturas cada paso sigue siendo rápido. exactitud compara la clase con el logit mayor, argmax(dim=1), con la etiqueta. Dentro de torch.no_grad(), autograd no registra las operaciones, porque medir no necesita gradientes.

perdida_fn = nn.CrossEntropyLoss()


def entrenar(modelo, optimizador, X, y, pasos=200):
    for _ in range(pasos):
        optimizador.zero_grad()
        perdida = perdida_fn(modelo(X), y)
        perdida.backward()
        optimizador.step()


def exactitud(modelo, X, y):
    with torch.no_grad():
        return (modelo(X).argmax(dim=1) == y).float().mean().item()

desde_cero entrena un Clasificador(4) nuevo con Adam y lr=0.01 sobre los ejemplos que recibe.

Pregunta. El modelo desde cero tiene 2820 parámetros y solo 20 ejemplos. ¿Qué exactitud esperas en esos 20 ejemplos? ¿Y en las 2000 muestras de prueba, comparada con 0.352?

def desde_cero(X, y):
    modelo = crear_modelo(4)
    entrenar(modelo, torch.optim.Adam(modelo.parameters(), lr=0.01), X, y)
    return modelo


modelo_cero = desde_cero(X_pocos, y_pocos)
print(round(exactitud(modelo_cero, X_pocos, y_pocos), 3))
print(round(exactitud(modelo_cero, X_prueba, y_destino_prueba), 3))

1.0
0.405

El preentrenamiento entrena el mismo tipo de red en la tarea fuente, con sus 20 000 ejemplos y una cabeza de 8 clases. Reutilizar después ese modelo en otra tarea se llama transfer learning. Seguir entrenando sus pesos con los datos de la tarea destino se llama fine-tuning o ajuste fino.

Pregunta. ¿El modelo fuente queda por encima o por debajo de 0.5 de exactitud en la prueba fuente? El azar con 8 clases balanceadas daría 0.125.

modelo_fuente = crear_modelo(8)
optimizador = torch.optim.Adam(modelo_fuente.parameters(), lr=0.01)
entrenar(modelo_fuente, optimizador, X_fuente, y_fuente)

print(round(exactitud(modelo_fuente, X_prueba, y_fuente_prueba), 3))

0.761

Pesos preentrenados

state_dict() devuelve un diccionario ordenado. Cada clave es el nombre de un tensor de parámetros y cada valor es el tensor. Los nombres siguen los atributos del modelo: cuerpo.0.weight es el peso de la primera capa de cuerpo. load_state_dict() carga los valores de un diccionario así en un modelo, clave por clave. La celda recorre en paralelo, con zip, los diccionarios de modelo_fuente y de un modelo nuevo de 4 clases, y después intenta la copia completa dentro de try.

Pregunta. ¿Qué pasa si intentas copiar el diccionario completo de modelo_fuente en un Clasificador(4) nuevo?

nuevo = crear_modelo(4)
for (nombre, fuente), destino in zip(
    modelo_fuente.state_dict().items(), nuevo.state_dict().values()
):
    print(f"{nombre:<16} {str(tuple(fuente.shape)):<9} {tuple(destino.shape)}")

try:
    nuevo.load_state_dict(modelo_fuente.state_dict())
except RuntimeError as error:
    print(str(error).splitlines()[1].strip())  # la línea que nombra la clave

cuerpo.0.weight  (32, 50)  (32, 50)
cuerpo.0.bias    (32,)     (32,)
cuerpo.2.weight  (32, 32)  (32, 32)
cuerpo.2.bias    (32,)     (32,)
cabeza.weight    (8, 32)   (4, 32)
cabeza.bias      (8,)      (4,)
size mismatch for cabeza.weight: copying a param with shape torch.Size([8, 32]) from checkpoint, the shape in current model is torch.Size([4, 32]).

La solución es copiar solo el cuerpo. copiar_cuerpo crea un modelo de 4 clases y carga en su atributo cuerpo el state_dict del cuerpo de origen.

Pregunta. Después de copiar_cuerpo(modelo_fuente), ¿comparten memoria el cuerpo copiado y el original?

def copiar_cuerpo(origen):
    modelo = crear_modelo(4)
    modelo.cuerpo.load_state_dict(origen.cuerpo.state_dict())
    return modelo


transferido = copiar_cuerpo(modelo_fuente)
cabeza_inicial = crear_modelo(4).cabeza.weight

print(torch.equal(transferido.cuerpo[0].weight, modelo_fuente.cuerpo[0].weight))
print(torch.equal(transferido.cabeza.weight, cabeza_inicial))
print(transferido.cuerpo[0].weight is modelo_fuente.cuerpo[0].weight)

True
True
False

Con modelos reales el flujo es el mismo. En la biblioteca Transformers de Hugging Face, AutoModelForSequenceClassification.from_pretrained("distilbert/distilbert-base-uncased", num_labels=2) descarga los pesos de DistilBERT, una versión reducida de BERT. El cuerpo es un transformer como el de Atención y transformers y recibe los pesos preentrenados. La cabeza de clasificación se crea con 2 salidas y pesos aleatorios, y la biblioteca avisa que esos pesos son nuevos y que el modelo necesita entrenarse en la tarea destino.

Cuerpo congelado

Congelar el cuerpo significa no cambiar sus pesos durante el entrenamiento en la tarea destino. requires_grad_(False) aplicado a un módulo marca todos sus parámetros para que autograd no les calcule gradiente. El optimizador recibe solo los parámetros de la cabeza. Así, entrenar solo ajusta la última capa sobre las 32 características que ya produce el cuerpo preentrenado. Para contar, p.numel() da el número de valores de cada tensor de parámetros.

Pregunta. ¿Cuántos de los 2820 parámetros quedan entrenables? ¿La exactitud en prueba queda por encima o por debajo de los 0.405 del modelo desde cero?

def congelar_cuerpo(origen, X, y):
    modelo = copiar_cuerpo(origen)
    modelo.cuerpo.requires_grad_(False)
    entrenar(modelo, torch.optim.Adam(modelo.cabeza.parameters(), lr=0.01), X, y)
    return modelo


congelado = congelar_cuerpo(modelo_fuente, X_pocos, y_pocos)
entrenables = sum(p.numel() for p in congelado.parameters() if p.requires_grad)
total = sum(p.numel() for p in congelado.parameters())

print(entrenables, total)
print(round(exactitud(congelado, X_prueba, y_destino_prueba), 3))
print(torch.equal(congelado.cuerpo[0].weight, modelo_fuente.cuerpo[0].weight))
print(torch.equal(congelado.cuerpo[2].weight, modelo_fuente.cuerpo[2].weight))

132 2820
0.649
True
True

Ajuste completo

El ajuste completo (full fine-tuning) entrena todos los parámetros: el cuerpo parte de los pesos preentrenados y también cambia con los datos destino. Un optimizador de PyTorch acepta una lista de grupos de parámetros. Cada grupo es un diccionario con sus "params" y, si hace falta, su propio "lr". Así el cuerpo y la cabeza pueden usar learning rates distintos.

def ajustar_todo(origen, X, y, lr_cuerpo):
    modelo = copiar_cuerpo(origen)
    optimizador = torch.optim.Adam(
        [
            {"params": modelo.cuerpo.parameters(), "lr": lr_cuerpo},
            {"params": modelo.cabeza.parameters(), "lr": 0.01},
        ]
    )
    entrenar(modelo, optimizador, X, y)
    return modelo

cambio_relativo mide cuánto se alejó el cuerpo de los pesos preentrenados. parameters_to_vector pone todos los parámetros del cuerpo en un solo vector. El resultado es la norma de la diferencia dividida entre la norma original: 0 significa que el cuerpo no cambió.

Pregunta. La cabeza usa lr=0.01 en los tres casos. Ordena lr_cuerpo = 0.01, 0.001 y 0.0001 según el cambio del cuerpo. ¿Alguno queda por debajo de los 0.649 del cuerpo congelado?

from torch.nn.utils import parameters_to_vector


def cambio_relativo(modelo, origen):
    nuevo = parameters_to_vector(modelo.cuerpo.parameters())
    antes = parameters_to_vector(origen.cuerpo.parameters())
    return ((nuevo - antes).norm() / antes.norm()).item()


for lr_cuerpo in (0.01, 0.001, 0.0001):
    ajustado = ajustar_todo(modelo_fuente, X_pocos, y_pocos, lr_cuerpo)
    print(
        f"lr_cuerpo={lr_cuerpo:<7}",
        f"cambio={cambio_relativo(ajustado, modelo_fuente):.3f}",
        f"exactitud={exactitud(ajustado, X_prueba, y_destino_prueba):.3f}",
    )

lr_cuerpo=0.01    cambio=0.497 exactitud=0.550
lr_cuerpo=0.001   cambio=0.116 exactitud=0.659
lr_cuerpo=0.0001  cambio=0.028 exactitud=0.659

Los modelos reales también se ajustan con learning rates bajos. La guía de clasificación de texto de Hugging Face ajusta todo DistilBERT, incluida la cabeza nueva, con un único learning_rate=2e-5, 50 veces menor que el lr_cuerpo=0.001 de esta lección, porque el cuerpo parte de un preentrenamiento que no conviene destruir.

Datos destino y tarea fuente

La ventaja del cuerpo preentrenado depende de cuántos ejemplos destino hay. exactitudes_medias entrena las tres estrategias con un número dado de ejemplos por clase. Repite el entrenamiento con tres subconjuntos distintos de ejemplos destino (semillas 10, 11 y 12) y devuelve la exactitud media en prueba de cada estrategia. El ajuste completo usa lr_cuerpo=0.001.

def exactitudes_medias(origen, por_clase):
    resultados = {"cero": [], "congelado": [], "completo": []}
    for semilla in (10, 11, 12):
        X, y = ejemplos_destino(por_clase, semilla)
        modelos = {
            "cero": desde_cero(X, y),
            "congelado": congelar_cuerpo(origen, X, y),
            "completo": ajustar_todo(origen, X, y, lr_cuerpo=0.001),
        }
        for nombre, modelo in modelos.items():
            resultados[nombre].append(exactitud(modelo, X_prueba, y_destino_prueba))
    return [np.mean(valores) for valores in resultados.values()]

comparar, en la celda plegada, llama a exactitudes_medias con cada tamaño e imprime una fila de tabla por tamaño con las tres medias.

def comparar(origen, tamanos):
    print("| por clase | cero | congelado | completo |")
    for por_clase in tamanos:
        cero, congelado, completo = exactitudes_medias(origen, por_clase)
        print(f"| {por_clase} | {cero:.3f} | {congelado:.3f} | {completo:.3f} |")

Pregunta. ¿Con qué tamaño es mayor la diferencia entre el cuerpo preentrenado y el entrenamiento desde cero? ¿La diferencia crece o se reduce con 100 ejemplos por clase?

La ventaja del cuerpo preentrenado también depende de la tarea fuente. pesos_fuente_parcial define una tarea fuente de 6 clases que solo usa los tres primeros factores. Con la semilla 1, generar produce las mismas lecturas de X_fuente; solo cambian las etiquetas. Para distinguir las 6 clases, el cuerpo no necesita el cuarto factor, pero la clase destino sí lo usa. La celda plegada entrena modelo_parcial con esas etiquetas, igual que modelo_fuente pero con una cabeza de 6 clases.

pesos_fuente_parcial = np.hstack([np.eye(4)[:, :3], -np.eye(4)[:, :3]])  # sin el cuarto factor
_, y_parcial, _ = generar(20000, semilla=1, pesos_clase_fuente=pesos_fuente_parcial)

modelo_parcial = crear_modelo(6)
optimizador = torch.optim.Adam(modelo_parcial.parameters(), lr=0.01)
entrenar(modelo_parcial, optimizador, X_fuente, y_parcial)

Pregunta. Con el cuerpo preentrenado en la tarea parcial, ¿qué estrategia gana con 5 ejemplos por clase? ¿Y con 100? ¿El cuerpo congelado alcanza al modelo desde cero?

Práctica

Traza guiada

Antes de procesar, escribe cuántos parámetros de traza quedan entrenables, cuántas claves conserva traza.state_dict() y si un paso de Adam sobre traza.parameters() cambiaría traza.cuerpo[0].weight y traza.cuerpo[2].weight.

traza = crear_modelo(4)
traza.cuerpo[0].requires_grad_(False)

Comprueba el conteo con sum(p.numel() for p in traza.parameters() if p.requires_grad) y las claves con len(traza.state_dict()). El conteo debe ser 2820 menos los parámetros de nn.Linear(50, 32). Si esperabas menos claves, recuerda que congelar no elimina parámetros: solo deja de calcular su gradiente.

Learning rate del cuerpo

Con ejemplos_destino(25, semilla=10), repite ajustar_todo para lr_cuerpo igual a 0.01, 0.003, 0.001, 0.0003, 0.0001 y 0.0. Antes de procesar, predice qué valor da la mejor exactitud en prueba y qué estrategia equivale a lr_cuerpo=0.0. Comprueba con assert que la exactitud con 0.0 es igual a la de congelar_cuerpo con los mismos ejemplos. Si no coincide, revisa que las dos funciones parten del mismo copiar_cuerpo(modelo_fuente).

Transferencia a otra tarea destino

Define una tarea destino de 3 clases con los mismos factores: pesos_otra = np.random.default_rng(5).normal(size=(4, 3)) y y = (factores @ pesos_otra).argmax(axis=1). Escribe un cálculo que:

1. defina generar_otra(n, semilla) a partir de generar, con pesos_otra en lugar de pesos_destino, y genere 10 ejemplos por clase para entrenar y 2000 muestras para prueba, con semillas distintas de 1 y 2
2. añada un parámetro n_clases a copiar_cuerpo para crear una cabeza de 3 salidas
3. entrene un modelo desde cero, uno con el cuerpo de modelo_fuente congelado y uno con ajuste completo y lr_cuerpo=0.001
4. compruebe con assert que el cuerpo congelado sigue igual al de modelo_fuente
5. escriba qué estrategia gana y relacione el resultado con los factores que usa la tarea fuente

Reto opcional

Programa un ajuste en dos fases: primero 100 pasos con el cuerpo congelado y después 100 pasos con todo el modelo y lr_cuerpo=0.001. Compáralo con ajustar_todo de 200 pasos, con modelo_fuente y con modelo_parcial, usando 5 ejemplos por clase. Mide cambio_relativo en ambos casos y explica qué gradientes recibe el cuerpo en los primeros pasos de cada versión.

Errores que revisar

• Si load_state_dict lanza size mismatch for cabeza.weight, estás copiando una cabeza de 8 clases en una de 4. Copia solo origen.cuerpo.state_dict() en modelo.cuerpo.
• Si el modelo transferido da la misma exactitud que el modelo desde cero, revisa que load_state_dict se ejecute antes de entrenar y sobre el mismo objeto que entrenas.
• Si modelo_fuente también cambia al ajustar el modelo destino, los dos comparten el mismo módulo, por ejemplo tras nuevo.cuerpo = modelo_fuente.cuerpo. Copia los pesos con load_state_dict.
• Si ajustar todo con pocos ejemplos queda muy por debajo de congelar y cambio_relativo es grande, como con lr_cuerpo=0.01, el cuerpo se está alejando de lo preentrenado. Dale al cuerpo su propio grupo en el optimizador con un learning rate menor que el de la cabeza.
• Si con muchos ejemplos el cuerpo congelado queda por debajo del modelo desde cero, al cuerpo le falta algo que la tarea destino necesita. Ajusta también el cuerpo.

Cierre

Explica por qué el cuerpo congelado de modelo_fuente supera al de modelo_parcial con 100 ejemplos por clase, aunque los dos se entrenaron con las mismas 20 000 lecturas.

Lista de comprobación

• ¿Qué parte del modelo cambia de forma cuando la tarea destino tiene otro número de clases?
• ¿Qué hace requires_grad_(False) con los gradientes y con los pesos del cuerpo?
• ¿Cómo se le da al cuerpo un learning rate distinto del de la cabeza?
• ¿Por qué un learning rate alto en el cuerpo puede dar peor resultado que congelarlo?
• ¿En qué situación supera el ajuste completo al cuerpo congelado?

Documentación oficial

• Transfer Learning for Computer Vision: ajuste completo y cuerpo congelado con una red preentrenada y una cabeza nueva.
• Saving and Loading Models: state_dict, load_state_dict y carga de pesos desde un modelo distinto.
• Autograd mechanics: requires_grad y cómo congelar partes de un modelo preentrenado.
• torch.optim: grupos de parámetros con learning rate propio.
• Text classification (Hugging Face): ajuste de DistilBERT con AutoModelForSequenceClassification, num_labels y Trainer.