Pipelines de batch processing en Python que no agotan la RAM

Cargar un dataset de entrenamiento entero con data = list(...) o np.array(...) funciona perfectamente hasta que deja de hacerlo: en cuanto el conjunto supera lo que cabe en RAM, el proceso termina con MemoryError o empieza a hacer swap hasta volverse inutilizable. Un generador de Python resuelve esto sin cambiar la forma en la que se escribe el bucle de entrenamiento: en vez de materializar todo el dataset de golpe, produce un lote a la vez, bajo demanda, y libera cada lote en cuanto se procesa. Este artículo construye ese pipeline pieza por pieza, desde el generador más simple hasta su integración con PyTorch y TensorFlow, y termina con un script de medición que confirma, con tracemalloc y sin dependencias externas, cuánta memoria se ahorra realmente frente a cargarlo todo de golpe.

El generador base: un lote a la vez, no el dataset entero

Un generador en Python es una función que, en lugar de devolver un valor y terminar, usa yield para producir una secuencia de valores uno a uno, suspendiendo su ejecución entre cada uno y conservando su estado local -variables, punto de ejecución- hasta la siguiente llamada, según la definición oficial de función generadora en la documentación de Python. Aplicado a batch processing, la promesa es que el proceso solo mantenga en memoria el lote actual, nunca el dataset completo -algo que depende de que la fuente de datos también sea perezosa, no solo el generador que la envuelve. Primero el mecanismo, con un array pequeño que ya vive entero en RAM (caso puramente didáctico, no el que importa):

from typing import Iterator
import numpy as np


def batch_generator(data: np.ndarray, batch_size: int) -> Iterator[np.ndarray]:
    """Produce lotes de `batch_size` elementos; NO evita cargar `data` si `data`
    ya es un array completo en memoria -ver la versión con archivo más abajo
    para el caso real de un dataset que no cabe en RAM."""
    total = len(data)
    for start in range(0, total, batch_size):
        yield data[start:start + batch_size]


dataset = np.arange(1_000_000)  # ejemplo pequeño, ya materializado en RAM
for batch in batch_generator(dataset, batch_size=256):
    pass  # aquí iría el paso de entrenamiento o inferencia sobre `batch`

La clave está en yield en vez de return: cada vuelta del bucle for reanuda la función justo donde se quedó, sin volver a ejecutar el cuerpo desde el principio ni acumular los lotes anteriores. Pero ese mecanismo por sí solo no cumple la promesa del título si data ya estaba entero en memoria antes de llamar a la función. La versión que sí la cumple lee de una fuente perezosa -un archivo línea a línea, por ejemplo-, de modo que en ningún momento el dataset completo llega a existir como objeto en memoria:

from pathlib import Path
from typing import Iterator, List


def file_batch_generator(path: str, batch_size: int) -> Iterator[List[str]]:
    """Lee el archivo línea a línea: nunca carga el archivo completo, solo
    el lote en curso más el búfer de lectura del sistema de archivos."""
    batch: List[str] = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            batch.append(line.rstrip("\n"))
            if len(batch) == batch_size:
                yield batch
                batch = []
    if batch:
        yield batch


if __name__ == "__main__":
    demo_path = Path("dataset.txt")
    demo_path.write_text("\n".join(f"muestra-{i}" for i in range(10_000)), encoding="utf-8")

    for batch in file_batch_generator(str(demo_path), batch_size=256):
        pass  # aquí iría el paso de entrenamiento o inferencia sobre `batch`

Esta es la pieza de base sobre la que se construye el resto del pipeline: el resto del artículo usa fuentes perezosas equivalentes -archivos, arrays mapeados en disco- precisamente porque son las que de verdad evitan cargar el dataset completo, no solo el generador que los envuelve.

Aleatoriedad y tipado: el generador que de verdad se usa en entrenamiento

El generador anterior recorre los datos siempre en el mismo orden, algo que en entrenamiento supervisado suele ser indeseable porque introduce sesgo de orden entre épocas. La versión que se usa en la práctica mezcla los índices antes de trocear, y para que el dataset siga sin caber-en-RAM lee features/labels con mmap_mode="r": el array queda mapeado en disco, y solo el lote indexado en cada vuelta llega a copiarse a memoria:

from typing import Iterator, Tuple
import numpy as np


def shuffled_batch_generator(
    features_path: str,
    labels_path: str,
    batch_size: int,
) -> Iterator[Tuple[np.ndarray, np.ndarray]]:
    features = np.load(features_path, mmap_mode="r")  # vive en disco, no en RAM
    labels = np.load(labels_path, mmap_mode="r")
    total = features.shape[0]
    order = np.random.permutation(total)
    for start in range(0, total, batch_size):
        idx = order[start:start + batch_size]
        yield features[idx], labels[idx]  # el indexado fancy sí copia, pero solo el lote


if __name__ == "__main__":
    np.save("features.npy", np.random.rand(1000, 10).astype(np.float32))
    np.save("labels.npy", np.random.randint(0, 2, size=1000).astype(np.int64))

    for x_batch, y_batch in shuffled_batch_generator("features.npy", "labels.npy", batch_size=64):
        pass  # entrenar_lote(x_batch, y_batch)

Con np.random.permutation el coste de memoria extra en RAM es el de un array de índices (un entero por muestra) más el lote en curso, no el del dataset completo. Es la pieza que se conecta directamente con los frameworks de ML en la siguiente sección: cualquier `Dataset`, `DataLoader` o `tf.data.Dataset` puede envolver este generador sin reescribirlo.

Cuando el framework espera un iterable: Dataset, DataLoader e IterableDataset en PyTorch

PyTorch define dos formas de dataset con contratos distintos, y elegir la incorrecta es la causa más común de código que "usa PyTorch" pero sigue cargando todo en memoria. Según la documentación oficial de torch.utils.data, un dataset map-style implementa __getitem__() y __len__() y representa un mapa de índices a muestras -acceso aleatorio-, mientras que un dataset iterable-style implementa __iter__() y representa un flujo de datos, pensado para casos donde la lectura aleatoria es cara o imposible, como bases de datos o streams en tiempo real.

Cuando el dataset admite acceso por índice, el patrón map-style delega el shuffle, el batching y el paralelismo por workers en DataLoader; para que el dataset siga sin caber en RAM, __getitem__ lee de un array mapeado en disco (mmap_mode="r") en vez de uno ya cargado, igual que en la sección anterior. El bloque de arranque va dentro de if __name__ == "__main__": porque, según la sección de comportamiento específico por plataforma de la documentación de torch.utils.data, en Windows y macOS spawn() es el método de arranque de workers por defecto: se lanza otro intérprete que vuelve a ejecutar el script principal antes de recibir el dataset. Sin ese guard, cada worker relanzaría el script entero de forma recursiva:

import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader


class ArrayDataset(Dataset):
    def __init__(self, features_path: str, labels_path: str) -> None:
        # mmap_mode="r": el archivo se mapea en disco, no se carga en RAM al abrirlo
        self.features = np.load(features_path, mmap_mode="r")
        self.labels = np.load(labels_path, mmap_mode="r")

    def __len__(self) -> int:
        return self.features.shape[0]

    def __getitem__(self, idx: int):
        x = torch.from_numpy(np.array(self.features[idx])).float()
        y = torch.tensor(int(self.labels[idx])).long()
        return x, y


if __name__ == "__main__":
    np.save("features.npy", np.random.rand(1000, 20).astype(np.float32))
    np.save("labels.npy", np.random.randint(0, 2, size=1000).astype(np.int64))

    dataset = ArrayDataset("features.npy", "labels.npy")
    loader = DataLoader(dataset, batch_size=64, shuffle=True, num_workers=2)

    for x_batch, y_batch in loader:
        pass  # entrenamiento

Cuando el origen ya es un flujo -por ejemplo, el generador con shuffle de la sección anterior-, IterableDataset evita forzar un índice que no existe. La propia documentación aclara que, con batch_size=None y sin batch_sampler, el automatic batching queda desactivado y DataLoader devuelve directamente lo que produce el iterador del dataset, que es exactamente lo que interesa cuando el generador ya entrega lotes. Dos detalles hay que resolver aquí y que un ejemplo ingenuo suele pasar por alto: guardar un generador ya creado en __init__ lo deja agotado después de la primera época (un generador solo se recorre una vez) y no es serializable para workers creados con spawn; y con num_workers > 0, la documentación de torch.utils.data.IterableDataset advierte que el mismo objeto dataset se replica en cada proceso worker, así que sin particionar explícitamente cada worker iteraría el generador completo y los datos quedarían duplicados tantas veces como workers activos. La solución a ambos problemas es la misma: no guardar un generador, sino algo capaz de crear uno nuevo cada vez que __iter__ lo pida, usando get_worker_info() para que cada worker reciba un rango de índices distinto. Ese "algo" tiene, además, una restricción propia de spawn (el método de arranque de workers por defecto en Windows y macOS): PyTorch serializa el dataset para enviarlo a cada proceso worker, y una función anidada dentro de otra función -una closure, como una fábrica que devuelve una función interna- no es serializable con pickle. La forma segura es un callable definido a nivel de módulo, por ejemplo una clase con __call__:

import math
import numpy as np
import torch
from torch.utils.data import IterableDataset, DataLoader, get_worker_info


class BatchSource:
    """Callable a nivel de módulo (no una función anidada/closure): para que
    DataLoader pueda serializarlo con pickle y enviarlo a cada worker cuando
    num_workers>0 usa spawn (Windows, macOS)."""
    def __init__(self, features_path: str, labels_path: str, batch_size: int) -> None:
        self.features_path = features_path
        self.labels_path = labels_path
        self.batch_size = batch_size

    def __call__(self, start: int, end: int):
        features = np.load(self.features_path, mmap_mode="r")
        labels = np.load(self.labels_path, mmap_mode="r")
        order = start + np.random.permutation(end - start)
        for s in range(0, len(order), self.batch_size):
            idx = order[s:s + self.batch_size]
            yield features[idx], labels[idx]


class StreamingDataset(IterableDataset):
    def __init__(self, batch_source: BatchSource, total_samples: int) -> None:
        # batch_source es una INSTANCIA de una clase de nivel de módulo, no una
        # closure: __iter__ puede invocarla de nuevo en cada época y en cada
        # worker (sin agotarse) y, a diferencia de una función anidada, sí es
        # picklable para spawn.
        self.batch_source = batch_source
        self.total_samples = total_samples

    def __iter__(self):
        worker_info = get_worker_info()
        if worker_info is None:
            start, end = 0, self.total_samples
        else:
            # Sin repartir el rango, cada worker iteraría el dataset completo
            # y los datos saldrían duplicados una vez por worker.
            per_worker = math.ceil(self.total_samples / worker_info.num_workers)
            start = worker_info.id * per_worker
            end = min(start + per_worker, self.total_samples)
        for x, y in self.batch_source(start, end):
            yield torch.from_numpy(np.array(x)).float(), torch.from_numpy(np.array(y)).long()


if __name__ == "__main__":
    TOTAL = 1000
    np.save("features.npy", np.random.rand(TOTAL, 20).astype(np.float32))
    np.save("labels.npy", np.random.randint(0, 2, size=TOTAL).astype(np.int64))

    batch_source = BatchSource("features.npy", "labels.npy", batch_size=64)
    dataset = StreamingDataset(batch_source, total_samples=TOTAL)
    loader = DataLoader(dataset, batch_size=None, num_workers=2)  # el generador ya entrega lotes completos

    for x_batch, y_batch in loader:
        pass  # entrenamiento

La regla práctica: si el dataset se puede indexar, usar Dataset + DataLoader en modo map-style y dejar que PyTorch gestione shuffle y workers; si el origen es un stream que ya decide su propio orden y tamaño de lote, envolverlo en IterableDataset en vez de forzarlo a parecer indexable.

La misma idea en TensorFlow: tf.data.Dataset.from_generator

TensorFlow expone el mismo patrón a través de tf.data.Dataset.from_generator, pero con un matiz que ha cambiado y que conviene verificar antes de copiar código de tutoriales antiguos: los parámetros output_types y output_shapes están marcados como deprecados en la documentación actual del método en tf.data.Dataset, que indica explícitamente usar output_signature en su lugar porque los argumentos antiguos se eliminarán en una versión futura. La forma vigente declara el tipo y la forma de cada elemento con tf.TensorSpec:

import tensorflow as tf
import numpy as np


def batch_gen():
    total = 10_000
    batch_size = 128
    for start in range(0, total, batch_size):
        end = min(start + batch_size, total)
        yield np.arange(start, end, dtype=np.int32)


dataset = tf.data.Dataset.from_generator(
    batch_gen,
    output_signature=tf.TensorSpec(shape=(None,), dtype=tf.int32),
)

for batch in dataset.take(3):
    print(batch.numpy())

El resultado sigue siendo un tf.data.Dataset normal, componible con map o filter como cualquier otro pipeline de tf.data -no con batch: cada elemento que produce batch_gen ya es un lote, así que volver a agrupar con .batch() anidaría lotes dentro de lotes en vez de aplanar nada-. El cambio principal frente a ejemplos antiguos que copien output_types/output_shapes es cómo se declara la forma de salida, no la lógica del generador en sí.

La propia documentación de from_generator añade una advertencia que conviene leer antes de adoptarlo por defecto: la implementación se apoya en tf.numpy_function y hereda sus mismas restricciones, lo que obliga a que las operaciones de dataset e iterador se ejecuten en el mismo proceso Python que llamó a from_generator -sujeto por tanto al GIL, sin poder escalar con tf.data service-, y el cuerpo del generador no se serializa en un GraphDef, así que no conviene usarlo si el modelo necesita guardarse y restaurarse en otro entorno.

Componer el pipeline en cadena: cargar, aumentar y agrupar sin acoplar

Hasta aquí cada generador hace una sola cosa: producir lotes. La ventaja real de los generadores en Python aparece al encadenarlos, porque cada etapa consume del generador anterior sin que ninguna materialice una lista intermedia completa:

from typing import Iterator, List, Tuple
import numpy as np


def load_samples(total: int) -> Iterator[Tuple[np.ndarray, int]]:
    """Genera cada muestra bajo demanda dentro del propio bucle; nunca
    construye la lista completa. En un pipeline real, el cuerpo del `for`
    leería una fila de un archivo o de una base de datos en vez de generar
    datos sintéticos."""
    for i in range(total):
        yield np.random.rand(10).astype(np.float32), i % 2


def augment(samples: Iterator[Tuple[np.ndarray, int]]) -> Iterator[Tuple[np.ndarray, int]]:
    for x, y in samples:
        # .astype(x.dtype): np.random.normal devuelve float64 por defecto y al
        # sumarlo a un x float32 lo promovería a float64 de forma silenciosa.
        noise = np.random.normal(0, 0.01, size=x.shape).astype(x.dtype)
        yield x + noise, y


def collect_batches(
    samples: Iterator[Tuple[np.ndarray, int]],
    batch_size: int,
) -> Iterator[List[Tuple[np.ndarray, int]]]:
    batch: List[Tuple[np.ndarray, int]] = []
    for sample in samples:
        batch.append(sample)
        if len(batch) == batch_size:
            yield batch
            batch = []
    if batch:
        yield batch


pipeline = collect_batches(augment(load_samples(200)), batch_size=32)
for batch in pipeline:
    inputs, targets = zip(*batch)
    # entrenamiento o inferencia sobre `inputs`, `targets`

Cada función se puede testear, sustituir o reordenar de forma independiente -cambiar augment por otra transformación no toca load_samples ni collect_batches-, y el pipeline completo sigue sin cargar nada que no sea el lote en curso.

Cuando la transformación necesita convertir a tensor y normalizar, tiene sentido envolver la cadena en una clase con estado propio en vez de una función suelta, para poder pasarle un transform configurable como se haría con un Dataset de PyTorch. Para ese transform, la API vigente de torchvision es torchvision.transforms.v2: la documentación oficial de transforms recomienda explícitamente usar los transforms v2 en vez de los de torchvision.transforms, porque son más rápidos y son los únicos que reciben nuevas funcionalidades a partir de ahora; de hecho, v2.ToTensor -habitual en tutoriales antiguos- figura como deprecado a favor de componer v2.ToImage() con v2.ToDtype(..., scale=True):

from typing import Callable, Generator, List, Optional, Tuple
import numpy as np
import torch
from torchvision.transforms import v2


class BatchGenerator:
    def __init__(
        self,
        dataset: List[Tuple[np.ndarray, int]],
        batch_size: int,
        transform: Optional[Callable[[torch.Tensor], torch.Tensor]] = None,
    ) -> None:
        self.dataset = dataset
        self.batch_size = batch_size
        self.transform = transform

    def __iter__(self) -> Generator[Tuple[torch.Tensor, torch.Tensor], None, None]:
        for start in range(0, len(self.dataset), self.batch_size):
            batch = self.dataset[start:start + self.batch_size]
            inputs, targets = zip(*batch)
            processed = [
                self.transform(torch.from_numpy(x)) if self.transform else torch.from_numpy(x)
                for x in inputs
            ]
            yield torch.stack(processed).float(), torch.tensor(targets).long()


# Los tensores de ejemplo ya son float32 en [0, 1], por eso basta con normalizar
# (si el origen fuera uint8 o PIL, ver la composición con ToImage/ToDtype más abajo).
transform = v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])

# `images` es una lista pequeña en RAM solo para que el ejemplo sea
# autocontenido; en producción `dataset` sería una secuencia indexable
# respaldada por disco -como el `ArrayDataset` con mmap de la sección
# anterior-, no una lista construida de golpe.
images = [(np.random.rand(3, 224, 224).astype(np.float32), i % 10) for i in range(500)]
batches = BatchGenerator(images, batch_size=32, transform=transform)

for x_batch, y_batch in batches:
    pass  # entrenamiento

Si el origen son imágenes PIL o arrays uint8 en vez de float32 ya normalizado, la composición recomendada por la documentación añade la conversión de tipo antes de normalizar: v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True), v2.Normalize(mean=..., std=...)]).

Medir el ahorro de memoria y elegir la pieza correcta

tracemalloc.get_traced_memory() devuelve una tupla (memoria_actual, pico) con el uso de memoria trazado por el asignador de Python en el momento exacto de la llamada, según la documentación oficial del módulo tracemalloc, sin necesidad de instrumentación externa ni de un profiler de terceros. El alcance importa: tracemalloc traza por defecto las asignaciones hechas directamente por el asignador de Python -el dominio 0, en los términos de su propia documentación-, es decir, objetos Python (listas, tuplas, enteros). No ve la memoria nativa que reservan por su cuenta buffers de NumPy, tensores de PyTorch o memoria de GPU; para esos casos hace falta memoria de proceso (resource.getrusage en Unix, un profiler como memory_profiler) o las utilidades de cada framework, como torch.cuda.memory_allocated(). El ejemplo siguiente usa listas de enteros de Python precisamente para que tracemalloc mida con exactitud lo que se le pide, sin ese matiz de por medio:

import tracemalloc


def load_all(n: int) -> list:
    return [i * i for i in range(n)]


def stream(n: int, batch_size: int):
    for start in range(0, n, batch_size):
        yield [i * i for i in range(start, min(start + batch_size, n))]


N = 5_000_000

tracemalloc.start()
data = load_all(N)
_, peak_full = tracemalloc.get_traced_memory()
tracemalloc.stop()
del data

tracemalloc.start()
for batch in stream(N, batch_size=10_000):
    pass
_, peak_stream = tracemalloc.get_traced_memory()
tracemalloc.stop()

print(f"Pico cargando todo en una lista: {peak_full / 1_048_576:.1f} MB")
print(f"Pico con generador por lotes:    {peak_stream / 1_048_576:.1f} MB")

Ejecutado tal cual en Python 3.13 sobre Windows, sin ninguna dependencia externa, este script imprime:

Pico cargando todo en una lista: 194.5 MB
Pico con generador por lotes:    0.8 MB

El resultado exacto depende de la máquina y de N, pero el patrón es el que importa: el pico de la versión con lista crece linealmente con N porque toda la secuencia vive en memoria a la vez; el pico de la versión con generador se mantiene aproximadamente constante porque solo existe un lote de batch_size elementos en cada momento. Ese es el resultado que hay que verificar en el dataset real antes de dar por buena cualquier de las piezas construidas arriba, porque un generador escrito de forma descuidada -por ejemplo, uno que hace list(iterable) antes de iterar, o que acumula resultados en una lista externa en vez de hacer yield de cada uno- puede acabar materializando exactamente lo que se quería evitar sin que el código deje de "parecer" un generador.

Ninguna de estas piezas es la respuesta universal: cuál usar depende de si el framework de destino ya define su propio contrato de datos (PyTorch, TensorFlow) o si el pipeline sigue siendo Python puro. Lo que no cambia entre las cuatro es el principio de fondo: medir con tracemalloc -o con el profiler de memoria que corresponda- antes de asumir que introducir generadores resolvió el problema, en vez de darlo por hecho porque el código ahora tiene un yield.