Diferencias entre Dataset e IterableDataset

Existen dos tipos de objetos de conjunto de datos: Dataset y IterableDataset. El tipo de conjunto de datos que elijas usar o crear depende del tamaño del conjunto de datos. En general, un IterableDataset es ideal para conjuntos de datos grandes (¡piensa en cientos de GB!) debido a su comportamiento perezoso y a sus ventajas de velocidad, mientras que un Dataset es excelente para todo lo demás. Esta página comparará las diferencias entre un Dataset y un IterableDataset para ayudarte a elegir el objeto de conjunto de datos adecuado para ti.

Descarga y transmisión

Cuando se dispone de un Dataset regular, se puede acceder a él mediante my_dataset[0]. Esto permite un acceso aleatorio a las filas. A este tipo de conjuntos de datos también se les llama conjuntos de datos “map-style”. Por ejemplo, se puede descargar ImageNet-1k de esta manera y acceder a cualquier fila:

from datasets import load_dataset

imagenet = load_dataset("timm/imagenet-1k-wds", split="train")  # downloads the full dataset
print(imagenet[0])

Pero hay una salvedad: debes tener el conjunto de datos completo almacenado en el disco o en memoria, lo que te impide acceder a conjuntos de datos más grandes que el disco. Como puede resultar incómodo para conjuntos de datos grandes, existe otro tipo de conjunto de datos, el IterableDataset. Cuando tienes un IterableDataset, puedes acceder a él usando un bucle for para cargar los datos de forma progresiva a medida que iteras sobre el conjunto de datos. De esta manera, solo se carga en memoria una pequeña fracción de ejemplos y no se escribe nada en el disco.

Por ejemplo, puede transmitir el conjunto de datos ImageNet-1k sin descargarlo en el disco:

from datasets import load_dataset

imagenet = load_dataset("timm/imagenet-1k-wds", split="train", streaming=True)  # will start loading the data when iterated over
for example in imagenet:
    print(example)
    break

Streaming puede leer datos en línea sin escribir ningún archivo en el disco. Por ejemplo, puedes transmitir conjuntos de datos compuestos por múltiples fragmentos, cada uno de cientos de gigabytes, como C4 o LAION-2B. Aprende más sobre cómo transmitir un conjunto de datos en la Guía de streaming de conjuntos de datos.

Sin embargo, esta no es la única diferencia, ya que el comportamiento “perezoso” de un IterableDataset también está presente cuando se trata de la creación y el procesamiento de conjuntos de datos.

Creación de datasets de estilo mapa y datasets iterables

Puedes crear un Dataset usando listas o diccionarios, y los datos se convierten por completo a Arrow para que puedas acceder fácilmente a cualquier fila:

my_dataset = Dataset.from_dict({"col_1": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]})
print(my_dataset[0])

Por otro lado, para crear un IterableDataset, debes proporcionar una forma “perezosa” de cargar los datos. En Python, generalmente se utilizan funciones generadoras. Estas funciones yield un ejemplo a la vez, lo que significa que no puedes acceder a una fila mediante rebanado como en un Dataset regular:

def my_generator(n):
    for i in range(n):
        yield {"col_1": i}

my_iterable_dataset = IterableDataset.from_generator(my_generator, gen_kwargs={"n": 10})
for example in my_iterable_dataset:
    print(example)
    break

Carga de archivos locales de forma completa y progresiva

Es posible convertir archivos de datos locales o remotos a un Dataset de Arrow usando load_dataset():

data_files = {"train": ["path/to/data.csv"]}
my_dataset = load_dataset("csv", data_files=data_files, split="train")
print(my_dataset[0])

Sin embargo, esto requiere un paso de conversión del formato CSV al formato Arrow, lo que toma tiempo y espacio en disco si tu conjunto de datos es grande.

Para ahorrar espacio en disco y omitir el paso de conversión, puede definir un IterableDataset transmitiendo directamente desde los archivos locales. De esta manera, los datos se leen progresivamente desde los archivos locales a medida que itera sobre el conjunto de datos:

data_files = {"train": ["path/to/data.csv"]}
my_iterable_dataset = load_dataset("csv", data_files=data_files, split="train", streaming=True)
for example in my_iterable_dataset:  # this reads the CSV file progressively as you iterate over the dataset
    print(example)
    break

Se admiten muchos formatos de archivo, como CSV, JSONL y Parquet, así como archivos de imagen y audio. Puede encontrar más información en las guías correspondientes para cargar conjuntos de datos tabulares, de texto, de visión y de audio.

Procesamiento de datos inmediato y procesamiento de datos diferido

Cuando se procesa un objeto Dataset con Dataset.map(), el conjunto de datos completo se procesa de inmediato y se devuelve. Esto es similar a cómo funciona, por ejemplo, pandas.

my_dataset = my_dataset.map(process_fn)  # process_fn is applied on all the examples of the dataset
print(my_dataset[0])

Por otro lado, debido a la naturaleza “perezosa” de un IterableDataset, la llamada a IterableDataset.map() no aplica tu función map sobre el conjunto de datos completo. En su lugar, tu función map se aplica sobre la marcha.

Por eso, puedes encadenar varios pasos de procesamiento y todos se ejecutarán a la vez cuando comiences a iterar sobre el conjunto de datos:

my_iterable_dataset = my_iterable_dataset.map(process_fn_1)
my_iterable_dataset = my_iterable_dataset.filter(filter_fn)
my_iterable_dataset = my_iterable_dataset.map(process_fn_2)

# process_fn_1, filter_fn and process_fn_2 are applied on-the-fly when iterating over the dataset
for example in my_iterable_dataset:  
    print(example)
    break

Barajado exacto y aproximado rápido

Cuando se baraja un Dataset usando Dataset.shuffle(), se aplica un barajado exacto del conjunto de datos. Funciona tomando una lista de índices [0, 1, 2, ... len(my_dataset) - 1] y barajando esta lista. Luego, acceder a my_dataset[0] devuelve la fila y el índice definidos por el primer elemento del mapeo de índices que ha sido barajado:

my_dataset = my_dataset.shuffle(seed=42)
print(my_dataset[0])

Dado que no tenemos acceso aleatorio a las filas en el caso de un IterableDataset, no podemos usar una lista de índices barajados y acceder a una fila en una posición arbitraria. Esto impide el uso de un barajado exacto. En su lugar, se utiliza un barajado aproximado rápido en IterableDataset.shuffle(). Utiliza un búfer de barajado para muestrear ejemplos aleatorios de forma iterativa del conjunto de datos. Dado que el conjunto de datos aún se lee de forma iterativa, ofrece un excelente rendimiento de velocidad:

my_iterable_dataset = my_iterable_dataset.shuffle(seed=42, buffer_size=100)
for example in my_iterable_dataset:
    print(example)
    break

Pero usar un buffer de mezcla no es suficiente para proporcionar una mezcla satisfactoria para el entrenamiento de modelos de aprendizaje automático. Por lo tanto, IterableDataset.shuffle() también mezcla los fragmentos del conjunto de datos si este está compuesto por varios archivos o fuentes. Además, rellena el buffer utilizando hasta max_buffer_input_shards fragmentos a la vez (el valor predeterminado es 10), lo que mejora considerablemente la calidad de la mezcla:

# Stream from the internet
my_iterable_dataset = load_dataset("deepmind/code_contests", split="train", streaming=True)
my_iterable_dataset.num_shards  # 39

# Stream from local files
data_files = {"train": [f"path/to/data_{i}.csv" for i in range(1024)]}
my_iterable_dataset = load_dataset("csv", data_files=data_files, split="train", streaming=True)
my_iterable_dataset.num_shards  # 1024

# From a generator function
def my_generator(n, sources):
    for source in sources:
        for example_id_for_current_source in range(n):
            yield {"example_id": f"{source}_{example_id_for_current_source}"}

gen_kwargs = {"n": 10, "sources": [f"path/to/data_{i}" for i in range(1024)]}
my_iterable_dataset = IterableDataset.from_generator(my_generator, gen_kwargs=gen_kwargs)
my_iterable_dataset.num_shards  # 1024

Si tu conjunto de datos tiene pocos archivos, puedes mejorar la calidad del barajado llamando a reshard() antes de barajar. Por ejemplo, para conjuntos de datos Parquet, esto vuelve a dividir los datos usando grupos de filas en lugar de tener un archivo por fragmento, lo que te proporciona muchos más fragmentos de los cuales barajar.

Diferencias de velocidad

Los objetos Dataset regulares se basan en Arrow, que proporciona un acceso aleatorio rápido a las filas. Gracias al mapeo de memoria y al hecho de que Arrow es un formato en memoria, la lectura de datos desde el disco no realiza llamadas de sistema costosas ni deserialización. Proporciona una carga de datos aún más rápida al iterar mediante un bucle for al iterar sobre lotes de registros de Arrow contiguos.

Sin embargo, en cuanto tu Dataset tiene un mapeo de índices (por ejemplo, mediante Dataset.shuffle()), la velocidad puede ser 10 veces más lenta. Esto se debe a que hay un paso adicional para obtener el índice de la fila que se debe leer usando el mapeo de índices y, lo más importante, ya no estás leyendo fragmentos de datos contiguos. Para restaurar la velocidad, necesitarías reescribir todo el conjunto de datos en tu disco de nuevo usando Dataset.flatten_indices(), lo cual elimina el mapeo de índices. Sin embargo, esto puede llevar mucho tiempo dependiendo del tamaño de tu conjunto de datos:

my_dataset[0]  # fast
my_dataset = my_dataset.shuffle(seed=42)
my_dataset[0]  # up to 10x slower
my_dataset = my_dataset.flatten_indices()  # rewrite the shuffled dataset on disk as contiguous chunks of data
my_dataset[0]  # fast again

En este caso, recomendamos cambiar a un IterableDataset y aprovechar su rápido método de barajado aproximado IterableDataset.shuffle(). Solo baraja el orden de los fragmentos y añade un búfer de barajado a tu conjunto de datos, lo que mantiene la velocidad de tu conjunto de datos óptima. También puedes volver a barajar el conjunto de datos fácilmente:

for example in enumerate(my_iterable_dataset):  # fast
    pass

shuffled_iterable_dataset = my_iterable_dataset.shuffle(seed=42, buffer_size=100)

for example in enumerate(shuffled_iterable_dataset):  # as fast as before
    pass

shuffled_iterable_dataset = my_iterable_dataset.shuffle(seed=1337, buffer_size=100)  # reshuffling using another seed is instantaneous

for example in enumerate(shuffled_iterable_dataset):  # still as fast as before
    pass

Si estás usando tu conjunto de datos en múltiples épocas, la semilla efectiva para barajar el orden de los fragmentos en el búfer de barajado es seed + epoch. Esto facilita volver a barajar un conjunto de datos entre épocas:

for epoch in range(n_epochs):
    my_iterable_dataset.set_epoch(epoch)
    for example in my_iterable_dataset:  # fast + reshuffled at each epoch using `effective_seed = seed + epoch`
        pass

Para reiniciar la iteración de un conjunto de datos de estilo mapa, simplemente puede omitir los primeros ejemplos:

my_dataset = my_dataset.select(range(start_index, len(dataset)))

Pero si usas un DataLoader con un Sampler, debes guardar el estado de tu sampler (es posible que hayas escrito un sampler personalizado que permita reanudar).

Por otro lado, los conjuntos de datos iterables no permiten el acceso aleatorio a un índice de ejemplo específico para reanudar desde ahí. Sin embargo, puedes usar IterableDataset.state_dict() y IterableDataset.load_state_dict() para reanudar desde un punto de control, de manera similar a lo que puedes hacer con los modelos y los optimizadores:

>>> iterable_dataset = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=3)
>>> # save in the middle of training
>>> state_dict = iterable_dataset.state_dict()
>>> # and resume later
>>> iterable_dataset.load_state_dict(state_dict)

Por debajo, el dataset iterable mantiene un registro del shard actual que se está leyendo y del índice del ejemplo en el shard actual, y guarda esta información en el state_dict.

Para reanudar desde un punto de control, el conjunto de datos omite todos los fragmentos que se leyeron previamente para reiniciar desde el fragmento actual. Luego, lee el fragmento y omite ejemplos hasta llegar al ejemplo exacto del punto de control.

Por lo tanto, reiniciar un conjunto de datos es bastante rápido, ya que no volverá a leer los fragmentos que ya se han iterado. Sin embargo, reanudar un conjunto de datos generalmente no es instantáneo, ya que debe volver a leer desde el principio del fragmento actual y omitir ejemplos hasta llegar a la ubicación del punto de control.

Esto se puede usar con StatefulDataLoader de torchdata, consulte streaming con un DataLoader de PyTorch.

Cambiar de estilo de mapa a iterable

Si desea aprovechar el comportamiento “perezoso” de un IterableDataset o sus ventajas de velocidad, puede cambiar su Dataset de estilo mapa a un IterableDataset:

my_iterable_dataset = my_dataset.to_iterable_dataset()

Si desea barajar su conjunto de datos o usarlo con un DataLoader de PyTorch, le recomendamos generar un IterableDataset fragmentado:

my_iterable_dataset = my_dataset.to_iterable_dataset(num_shards=1024)
my_iterable_dataset.num_shards  # 1024