Procesar grandes volúmenes de texto en memoria es uno de los cuellos de botella más comunes al entrenar modelos de NLP. Crear un DataLoader en PyTorch junto con una collate function te permite dividir el dataset en lotes pequeños, optimizar el uso de GPU y preparar tus datos para entrenar modelos de clasificación como los que se construyen sobre DBpedia.
Qué hace un DataLoader y por qué lo necesitas
Un DataLoader es una utilidad de torch.utils.data que toma un dataset completo y lo divide en batches, es decir, lotes pequeños que el modelo procesa de forma iterativa. Esto te permite tres cosas concretas: ahorrar memoria, paralelizar el cómputo entre varias GPUs y mezclar datos antes de cada época.
¿Qué es un batch en PyTorch? Es un grupo pequeño de ejemplos que el modelo procesa al mismo tiempo. En vez de cargar miles de textos a la vez, defines un batch_size (por ejemplo, 8) y el DataLoader te entrega los datos en porciones manejables.
En la implementación que verás más abajo, el dataset es el iterador de entrenamiento de DBpedia y el batch_size elegido es 8, un tamaño adecuado para una GPU de Colab [01:30]. Si tienes más poder de cómputo, puedes subir a 16, 32 o 64.
Cómo elegir el batch size correcto
El tamaño del lote depende del hardware disponible y de la complejidad del dataset:
- 8: ideal para GPUs modestas como las de Colab gratuito.
- 16 o 32: buen balance entre velocidad y memoria en GPUs intermedias.
- 64 o más: requiere GPUs potentes y datasets bien optimizados.
DBpedia no incluye múltiples idiomas ni estructuras complejas, así que un batch pequeño funciona sin problema [01:50].
Por qué necesitas una collate function antes del DataLoader
La collate function es la pieza que une los ejemplos individuales del batch en un solo tensor. Sin ella, el DataLoader no sabe cómo combinar textos de longitudes distintas en una estructura uniforme. Por eso, antes de instanciar el DataLoader debes definir esta función y pasarla como argumento.
¿Qué es una collate function? Es una función que recibe un grupo de ejemplos (textos y etiquetas) y los concatena en tensores listos para alimentar al modelo. Convierte listas de Python en tensores de PyTorch.
Cómo configurar el device para aprovechar la GPU
Lo primero dentro de la función es decidir dónde se va a procesar la información. Si CUDA está disponible, usas GPU; si no, CPU [03:20]:
python
import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
Esto garantiza que tus tensores vivan en el dispositivo correcto y evita errores cuando el modelo espera datos en GPU pero los recibe en CPU.
Cómo construir paso a paso la función collate_batch
La función collate_batch recibe un lote de tuplas (etiqueta, texto) y devuelve tres tensores: etiquetas, textos concatenados y offsets. Empiezas inicializando tres listas:
- label_list: vacía, guardará las etiquetas procesadas.
- text_list: vacía, guardará cada texto convertido en tensor.
- offsets: inicia con el valor 0 y registrará dónde comienza cada texto dentro del tensor concatenado.
Dentro del bucle for, recorres cada par (label, text) del batch. Aplicas el label_pipeline a la etiqueta y el text_pipeline al texto, este último convertido a un tensor de tipo torch.int64 porque el vocabulario solo devuelve enteros y así ahorras memoria [05:40].
python
processed_text = torch.tensor(text_pipeline(text), dtype=torch.int64)
text_list.append(processed_text)
offsets.append(processed_text.size(0))
Qué son los offsets y por qué importan
Los offsets son posiciones que indican dónde empieza cada texto dentro del tensor grande que vas a concatenar. Si el primer texto mide 10 tokens, el siguiente empieza en la posición 10. Si el segundo mide 5, el tercero empieza en la posición 15.
¿Para qué sirven los offsets en NLP? Marcan el inicio de cada secuencia cuando concatenas varios textos en un solo tensor. Sin offsets, el modelo no sabría dónde termina un ejemplo y empieza otro.
Para generar los offsets correctos, aplicas la función cumsum (suma acumulativa) sobre la dimensión 0:
python
offsets = torch.tensor(offsets[:-1]).cumsum(dim=0)
Eliminas el último elemento porque no necesitas marcar el inicio de un texto que no existe, y cumsum convierte la lista de tamaños en posiciones acumuladas. La dimensión 0 en PyTorch siempre se refiere a las filas [07:50].
Cómo concatenar y devolver los tensores finales
El último paso es convertir la lista de etiquetas en tensor, concatenar todos los textos en uno solo y enviar todo al device:
python
label_list = torch.tensor(label_list, dtype=torch.int64)
text_list = torch.cat(text_list)
return label_list.to(device), text_list.to(device), offsets.to(device)
La función torch.cat pega todos los tensores de texto en uno solo, y los offsets te permiten reconstruir las fronteras entre ejemplos cuando el modelo los necesite.
Cómo instanciar finalmente el DataLoader
Una vez tienes la collate function lista, creas el DataLoader pasando el dataset, el batch_size, la opción de shuffle y la función de recolección:
python
from torch.utils.data import DataLoader
train_iter = DBpedia(split='train')
dataloader = DataLoader(train_iter, batch_size=8, shuffle=False, collate_fn=collate_batch)
En este caso shuffle queda en False porque DBpedia ya viene mezclado previamente [02:30]. Si trabajas con datasets ordenados por clase, activa shuffle=True para evitar sesgos durante el entrenamiento.
Con esto tu pipeline de datos queda listo para alimentar un modelo de clasificación de texto. ¿Te ha quedado clara la lógica de los offsets? Cuéntame en los comentarios qué batch_size usas según tu GPU.