Transfer Learning con TensorFlow y PyTorch: Más Allá de la Congelación de Capas

El Transfer Learning (Aprendizaje por Transferencia) es una de las técnicas más poderosas y ampliamente utilizadas en el campo de la Inteligencia Artificial, especialmente en Deep Learning. Permite aprovechar el conocimiento adquirido por un modelo entrenado en una tarea y conjunto de datos grandes (por ejemplo, clasificación de imágenes en ImageNet) para resolver una nueva tarea relacionada con menos datos y recursos computacionales. En lugar de entrenar un modelo desde cero, se toma un modelo ya existente y se lo adapta a nuestra necesidad específica. Esto es particularmente útil en escenarios donde la recopilación de grandes conjuntos de datos es costosa o imposible.

En este tutorial, exploraremos a fondo el Transfer Learning, desde sus fundamentos hasta técnicas avanzadas, y cómo implementarlo eficazmente utilizando las dos bibliotecas de Deep Learning más populares: TensorFlow y PyTorch.

💡 ¿Por Qué Usar Transfer Learning?

Imagínate que quieres construir un clasificador de imágenes que distinga entre diferentes razas de perros. Entrenar un modelo de red neuronal convolucional (CNN) desde cero para esta tarea requeriría una enorme cantidad de imágenes de perros etiquetadas y una considerable potencia computacional. Sin embargo, un modelo como ResNet o VGG, ya entrenado en el gigantesco conjunto de datos ImageNet (que contiene millones de imágenes de 1000 categorías diferentes), ya ha aprendido a detectar características de bajo nivel (bordes, texturas) y de alto nivel (formas, objetos).

Las principales razones para adoptar Transfer Learning son:

• Reducción de la Necesidad de Datos: Requiere menos datos etiquetados para la nueva tarea.
• Aceleración del Entrenamiento: El modelo ya tiene una base sólida, lo que reduce el tiempo de entrenamiento.
• Mejora del Rendimiento: A menudo, los modelos resultantes tienen un mejor rendimiento, especialmente con conjuntos de datos pequeños.
• Menos Recursos Computacionales: Se necesita menos hardware potente para el entrenamiento.

📖 Fundamentos del Transfer Learning

El Transfer Learning se basa en la idea de que las características aprendidas por las capas iniciales de una red neuronal convolucional (CNN) son genéricas y pueden ser útiles para diversas tareas de visión por computadora. Las capas posteriores, en cambio, aprenden características más específicas de la tarea para la que fue entrenado originalmente el modelo.

Tipos de Transfer Learning

Existen varias estrategias para aplicar Transfer Learning:

1. Congelación de Capas (Feature Extraction): Se toman las capas convolucionales de un modelo pre-entrenado y se 'congelan', es decir, sus pesos no se actualizan durante el entrenamiento. Se añaden nuevas capas (generalmente una o más capas densas) al final del modelo y solo estas nuevas capas se entrenan. Esto convierte el modelo pre-entrenado en un extractor de características.
2. Ajuste Fino (Fine-tuning): Después de la congelación inicial, o directamente, se 'descongelan' algunas o todas las capas del modelo pre-entrenado y se entrenan junto con las nuevas capas añadidas. Esto permite que el modelo ajuste los pesos de las capas pre-entrenadas para adaptarlas mejor a la nueva tarea. A menudo, se utiliza una tasa de aprendizaje más pequeña para las capas pre-entrenadas para evitar 'olvidar' el conocimiento original.
3. Entrenamiento desde Cero con Pesos Pre-entrenados: Se inicializa un modelo con los pesos pre-entrenados y luego se entrena todo el modelo (todas las capas) en el nuevo conjunto de datos. Esto es similar al ajuste fino, pero a menudo se aplica cuando el nuevo conjunto de datos es muy grande y similar al original, y queremos que el modelo aprenda la nueva tarea de manera más completa.

🛠️ Implementación en TensorFlow

Vamos a ver cómo implementar Transfer Learning en TensorFlow usando Keras. Utilizaremos el modelo pre-entrenado MobileNetV2 para clasificar un pequeño conjunto de datos de imágenes de flores.

1. Preparación de Datos

Primero, necesitamos un conjunto de datos. Usaremos el conjunto de datos de flores de TensorFlow.

import tensorflow as tf
import tensorflow_datasets as tfds

IMG_SIZE = (160, 160)
BATCH_SIZE = 32

# Cargar el dataset de flores
dataset, info = tfds.load(
    'tf_flowers',
    with_info=True,
    as_supervised=True,
    split=['train[:80%]', 'train[80%:]'],
)

train_dataset = dataset[0]
validation_dataset = dataset[1]

def preprocess(image, label):
    image = tf.image.resize(image, IMG_SIZE)
    image = tf.keras.applications.mobilenet_v2.preprocess_input(image)
    return image, label

train_dataset = train_dataset.map(preprocess).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)
validation_dataset = validation_dataset.map(preprocess).batch(BATCH_SIZE).prefetch(tf.data.AUTOTUNE)

2. Carga del Modelo Pre-entrenado y Congelación de Capas

Cargaremos MobileNetV2 pre-entrenado en ImageNet, sin incluir las capas clasificadoras superiores (include_top=False).

base_model = tf.keras.applications.MobileNetV2(input_shape=IMG_SIZE + (3,),
                                               include_top=False,
                                               weights='imagenet')

base_model.trainable = False # Congelar las capas del modelo base

3. Añadir Capas Clasificadoras y Entrenar

Añadimos una capa de pooling y una capa densa para la clasificación de nuestras flores.

num_classes = info.features['label'].num_classes # Número de clases de flores

model = tf.keras.Sequential([
    base_model,
    tf.keras.layers.GlobalAveragePooling2D(),
    tf.keras.layers.Dense(num_classes, activation='softmax')
])

model.compile(optimizer=tf.keras.optimizers.Adam(),
              loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])

# Entrenar solo las nuevas capas
history_feature_extraction = model.fit(train_dataset,
                                         epochs=10,
                                         validation_data=validation_dataset)

4. Ajuste Fino (Fine-tuning)

Para el ajuste fino, primero 'descongelamos' el modelo base o algunas de sus capas y luego lo entrenamos con una tasa de aprendizaje más baja.

base_model.trainable = True # Descongelar el modelo base

# Congelar capas inferiores para evitar la sobreescritura temprana de pesos
for layer in base_model.layers[:-50]: # Descongela las últimas 50 capas por ejemplo
    layer.trainable = False

model.compile(optimizer=tf.keras.optimizers.Adam(learning_rate=0.0001), # Tasa de aprendizaje más baja
              loss=tf.keras.losses.SparseCategoricalCrossentropy(),
              metrics=['accuracy'])

# Continuar el entrenamiento con ajuste fino
history_fine_tune = model.fit(train_dataset,
                               epochs=10, # Más epochs para ajuste fino
                               validation_data=validation_dataset,
                               initial_epoch=history_feature_extraction.epoch[-1])

🚀 Implementación en PyTorch

Ahora, veamos cómo lograr lo mismo utilizando PyTorch. Reutilizaremos un modelo ResNet pre-entrenado.

1. Preparación de Datos

PyTorch tiene su propia forma de manejar los datasets y dataloaders. Usaremos torchvision.datasets.

import torch
import torch.nn as nn
import torch.optim as optim
from torchvision import datasets, models, transforms
from torch.utils.data import DataLoader
import os

# Rutas a tus datos de flores (asumiendo que los tienes organizados en carpetas train/val/clase)
data_dir = 'path/to/your/flower_photos'

# Transformaciones para preprocesamiento
data_transforms = {
    'train': transforms.Compose([
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
    'val': transforms.Compose([
        transforms.Resize(256),
        transforms.CenterCrop(224),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ]),
}

# Cargar datasets
image_datasets = {
    x: datasets.ImageFolder(os.path.join(data_dir, x), data_transforms[x])
    for x in ['train', 'val']
}

dataloaders = {
    x: DataLoader(image_datasets[x], batch_size=4, shuffle=True, num_workers=4)
    for x in ['train', 'val']
}

dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

2. Carga del Modelo Pre-entrenado y Congelación de Capas

Cargamos ResNet18 pre-entrenado y congelamos todos sus parámetros.

model_ft = models.resnet18(pretrained=True)

for param in model_ft.parameters():
    param.requires_grad = False # Congelar todos los parámetros

# Modificar la última capa completamente conectada (fc)
num_ftrs = model_ft.fc.in_features
model_ft.fc = nn.Linear(num_ftrs, len(class_names))

model_ft = model_ft.to(device)

# Solo se optimizan los parámetros de la nueva capa fc
optimizer_ft = optim.Adam(model_ft.fc.parameters(), lr=0.001)
criterion = nn.CrossEntropyLoss()

# Función de entrenamiento simplificada (realmente se necesita un bucle completo)
# for epoch in range(num_epochs):
#     for inputs, labels in dataloaders['train']:
#         inputs = inputs.to(device)
#         labels = labels.to(device)
#         optimizer_ft.zero_grad()
#         outputs = model_ft(inputs)
#         loss = criterion(outputs, labels)
#         loss.backward()
#         optimizer_ft.step()
# ... (código de evaluación)

3. Ajuste Fino (Fine-tuning)

Para el ajuste fino, 'descongelamos' algunas capas del modelo base y usamos una tasa de aprendizaje más pequeña para ellas.

# Descongelar las últimas capas, por ejemplo, todas las capas de ResNet
for param in model_ft.parameters():
    param.requires_grad = True # Descongelar todos los parámetros para ajuste fino

# Opcional: Establecer tasas de aprendizaje diferentes para capas diferentes
# Por ejemplo, una tasa de aprendizaje más baja para las capas convolucionales y normal para la nueva capa fc
optimizer_ft = optim.SGD(model_ft.parameters(), lr=0.0001, momentum=0.9)
# También se puede usar optim.Adam(model_ft.parameters(), lr=0.00001)

# ... (bucle de entrenamiento con el nuevo optimizador)

📈 Estrategias Avanzadas de Transfer Learning

El Transfer Learning no se limita a congelar o descongelar capas. Aquí hay algunas técnicas avanzadas:

1. Descongelación Gradual (Layer-wise Unfreezing)

En lugar de descongelar todo a la vez, se puede descongelar el modelo por fases, comenzando por las capas más cercanas a la salida y progresando hacia las capas de entrada. Esto permite que el modelo se adapte gradualmente.

2. Detección de Parámetros Específicos para Ajuste Fino

Algunos modelos pre-entrenados tienen nombres de capas muy específicos. Puedes inspeccionar la estructura del modelo y decidir qué bloques o capas específicos quieres descongelar.

Por ejemplo, en PyTorch, puedes ver los nombres de las capas:

# print(model_ft)
# for name, param in model_ft.named_parameters():
#     print(name, param.requires_grad)

Esto te permitiría hacer cosas como: if 'layer4' in name: param.requires_grad = True para descongelar solo la última etapa de convoluciones.

3. Adapters o LoRA (Low-Rank Adaptation)

Para tareas muy específicas o para reducir el número de parámetros entrenables, se pueden insertar pequeñas redes neuronales (adaptadores) en puntos estratégicos del modelo pre-entrenado, en lugar de modificar las capas existentes. LoRA es una técnica popular que inyecta matrices de bajo rango en las matrices de peso de los transformadores, reduciendo drásticamente el número de parámetros entrenables mientras se mantienen la calidad del modelo original.

4. Congelación por Lote (Batch Normalization Freezing)

Las capas de Batch Normalization (BN) se comportan de manera diferente durante el entrenamiento y la inferencia. Si congelas capas con BN, a menudo querrás mantener las capas BN en modo de evaluación (eval() en PyTorch, training=False en TensorFlow) para que utilicen las estadísticas medias/varianza aprendidas durante el pre-entrenamiento, en lugar de calcular nuevas estadísticas con pequeños lotes de tu nuevo dataset, lo que podría desestabilizar el entrenamiento.

📊 Métricas y Evaluación del Rendimiento

Después de aplicar Transfer Learning, es crucial evaluar el rendimiento de tu modelo. Las métricas comunes incluyen:

• Precisión (Accuracy): Proporción de predicciones correctas.
• Precisión (Precision): De todas las predicciones positivas, cuántas fueron correctas.
• Recall: De todos los casos positivos reales, cuántos fueron predichos correctamente.
• F1-Score: Media armónica de precisión y recall.
• Matriz de Confusión: Tabla que resume el rendimiento del clasificador.

Es importante monitorear las curvas de pérdida y precisión tanto en el conjunto de entrenamiento como en el de validación para detectar sobreajuste (overfitting) o infraajuste (underfitting).

✅ Buenas Prácticas y Consejos

• Empieza Simple: Comienza congelando todo el modelo base y solo entrena las nuevas capas. Si el rendimiento no es suficiente, procede con el ajuste fino.
• Dataset Similaridad: Cuanto más similar sea tu conjunto de datos al conjunto de datos original con el que se entrenó el modelo base (e.g., ImageNet), más efectivas serán las capas pre-entrenadas.
• Tamaño del Dataset:
  • Pocos datos, similar al original: Congelación de características. Entrenar solo las últimas capas.
  • Pocos datos, diferente al original: Podrías necesitar un ajuste fino más profundo o incluso un pre-entrenamiento si tienes datos intermedios.
  • Muchos datos, similar al original: Ajuste fino de todas las capas con una LR baja.
  • Muchos datos, diferente al original: Podrías entrenar desde cero o hacer un ajuste fino muy agresivo si el modelo pre-entrenado es un buen punto de partida.
• Tasa de Aprendizaje: Usa tasas de aprendizaje bajas para el ajuste fino. Considera usar programadores de tasa de aprendizaje (learning rate schedulers) que disminuyan la tasa de aprendizaje con el tiempo.
• Regularización: El Dropout en las nuevas capas añadidas y la regularización L2 pueden ayudar a prevenir el sobreajuste durante el ajuste fino.
• Normalización de Entrada: Asegúrate de que tus datos de entrada sean preprocesados de la misma manera que lo fueron para el modelo pre-entrenado (normalización, tamaño de imagen, etc.). Los modelos de tf.keras.applications y torchvision.models suelen tener funciones específicas de preprocesamiento.
• Explora Diferentes Arquitecturas: No te limites a una sola. Prueba diferentes modelos pre-entrenados (ResNet, VGG, EfficientNet, Vision Transformers, etc.) para ver cuál se adapta mejor a tu tarea.

❓ Preguntas Frecuentes (FAQs)

Conclusión ✨

El Transfer Learning es una piedra angular en el Deep Learning moderno, permitiendo a desarrolladores y científicos de datos construir modelos de alto rendimiento con menos recursos. Dominar la congelación de capas y el ajuste fino, tanto en TensorFlow como en PyTorch, es una habilidad invaluable. Al aplicar estas técnicas y seguir las mejores prácticas, puedes acelerar significativamente el desarrollo de tus aplicaciones de IA y obtener resultados impresionantes, incluso con recursos limitados. ¡Ahora estás listo para llevar tus modelos de Deep Learning al siguiente nivel!