Diferenças entre Dataset e IterableDataset

Existem dois tipos de objetos de dataset, o Dataset e o IterableDataset. O tipo de dataset que você escolher usar ou criar depende do tamanho do dataset. Em geral, um IterableDataset é ideal para datasets grandes (pense em centenas de GBs!) devido ao seu comportamento preguiçoso e às vantagens de velocidade, enquanto um Dataset é ótimo para tudo o mais. Esta página comparará as diferenças entre um Dataset e um IterableDataset para ajudar você a escolher o objeto de dataset certo para você.

Baixando e transmitindo

Quando você tem um Dataset regular, pode acessá-lo usando my_dataset[0]. Isso fornece acesso aleatório às linhas. Esses conjuntos de dados também são chamados de conjuntos de dados “map-style”. Por exemplo, você pode baixar o ImageNet-1k da seguinte forma e acessar qualquer linha:

from datasets import load_dataset

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

No entanto, uma ressalva é que você deve ter o conjunto de dados inteiro armazenado no disco ou na memória, o que impede o acesso a conjuntos de dados maiores que o disco. Como isso pode se tornar inconveniente para conjuntos de dados grandes, existe outro tipo de conjunto de dados, o IterableDataset. Quando você tem um IterableDataset, pode acessá-lo usando um loop for para carregar os dados progressivamente à medida que itera sobre o conjunto de dados. Dessa forma, apenas uma pequena fração dos exemplos é carregada na memória, e você não grava nada no disco.

Por exemplo, você pode transmitir o conjunto de dados ImageNet-1k sem baixá-lo no 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

O Streaming pode ler dados online sem gravar nenhum arquivo no disco. Por exemplo, você pode transmitir datasets compostos por vários shards, cada um com centenas de gigabytes, como C4 ou LAION-2B. Saiba mais sobre como transmitir um dataset no Dataset Streaming Guide.

No entanto, essa não é a única diferença, pois o comportamento “lazy” de um IterableDataset também está presente quando se trata da criação e do processamento do dataset.

Criando datasets no estilo mapa e datasets iteráveis

Você pode criar um Dataset usando listas ou dicionários, e os dados são totalmente convertidos para Arrow, permitindo que você acesse facilmente qualquer linha:

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

Por outro lado, para criar um IterableDataset, você deve fornecer uma maneira “preguiçosa” de carregar os dados. Em Python, geralmente usamos funções geradoras. Essas funções yield um exemplo de cada vez, o que significa que você não pode acessar uma linha por fatiamento, como em um 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

Carregando arquivos locais de forma integral e progressiva

É possível converter arquivos de dados locais ou remotos em um Dataset 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])

No entanto, isso requer uma etapa de conversão do formato CSV para o formato Arrow, o que leva tempo e ocupa espaço em disco se o seu conjunto de dados for grande.

Para economizar espaço em disco e pular a etapa de conversão, você pode definir um IterableDataset fazendo streaming diretamente dos arquivos locais. Dessa forma, os dados são lidos progressivamente dos arquivos locais à medida que você itera sobre o dataset:

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

Muitos formatos de arquivo são suportados, como CSV, JSONL e Parquet, além de arquivos de imagem e áudio. Você pode encontrar mais informações nos guias correspondentes para carregar conjuntos de dados tabulares, de texto, de visão e de áudio.

Processamento de dados imediato e processamento de dados preguiçoso

Ao processar um objeto Dataset usando Dataset.map(), o conjunto de dados inteiro é processado imediatamente e retornado. Isso é semelhante à forma como o pandas funciona, por exemplo.

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

Por outro lado, devido à natureza “preguiçosa” de um IterableDataset, a chamada de IterableDataset.map() não aplica sua função map sobre o conjunto de dados completo. Em vez disso, sua função map é aplicada em tempo real.

Por isso, você pode encadear múltiplas etapas de processamento e todas elas serão executadas de uma vez quando você começar a iterar sobre o conjunto de dados:

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

Embaralhamento exato e aproximado rápido

Ao embaralhar um Dataset usando Dataset.shuffle(), você aplica um embaralhamento exato do conjunto de dados. Isso funciona obtendo uma lista de índices [0, 1, 2, ... len(my_dataset) - 1] e embaralhando essa lista. Em seguida, acessar my_dataset[0] retorna a linha e o índice definidos pelo primeiro elemento do mapeamento de índices que foi embaralhado:

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

Como não temos acesso aleatório às linhas no caso de um IterableDataset, não podemos usar uma lista embaralhada de índices e acessar uma linha em uma posição arbitrária. Isso impede o uso de embaralhamento exato. Em vez disso, um embaralhamento aproximado rápido é usado em IterableDataset.shuffle(). Ele usa um buffer de embaralhamento para amostrar exemplos aleatórios iterativamente do conjunto de dados. Como o conjunto de dados ainda é lido iterativamente, ele oferece excelente desempenho de velocidade:

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

Mas usar um buffer de embaralhamento não é suficiente para fornecer um embaralhamento satisfatório para o treinamento de modelos de aprendizado de máquina. Portanto, IterableDataset.shuffle() também embaralha os shards do conjunto de dados, caso seu conjunto de dados seja composto por múltiplos arquivos ou fontes. Além disso, ele preenche o buffer usando até max_buffer_input_shards shards por vez (o padrão é 10), o que melhora significativamente a qualidade do embaralhamento:

# 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

Se o seu conjunto de dados tiver poucos arquivos, você pode melhorar a qualidade do embaralhamento chamando reshard() antes de embaralhar. Por exemplo, para conjuntos de dados Parquet, isso reorganiza os shards usando grupos de linhas em vez de ter um arquivo por shard, fornecendo muitos mais shards para embaralhar.

Diferenças de velocidade

Objetos Dataset regulares são baseados em Arrow, que fornece acesso aleatório rápido às linhas. Graças ao mapeamento de memória e ao fato de Arrow ser um formato em memória, a leitura de dados do disco não realiza chamadas de sistema caras nem desserialização. Ele fornece carregamento de dados ainda mais rápido ao iterar usando um loop for, iterando sobre lotes de registros Arrow contíguos.

No entanto, assim que o seu Dataset tiver um mapeamento de índices (por exemplo, via Dataset.shuffle()), a velocidade pode ficar 10x mais lenta. Isso ocorre porque há uma etapa extra para obter o índice da linha a ser lido usando o mapeamento de índices e, mais importante, você não está mais lendo blocos contíguos de dados. Para restaurar a velocidade, você precisaria reescrever o conjunto de dados inteiro novamente no seu disco usando Dataset.flatten_indices(), que remove o mapeamento de índices. No entanto, isso pode levar muito tempo, dependendo do tamanho do seu conjunto de dados:

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

Nesse caso, recomendamos alternar para um IterableDataset e aproveitar seu método rápido de embaralhamento aproximado IterableDataset.shuffle(). Ele apenas embaralha a ordem dos shards e adiciona um buffer de embaralhamento ao seu dataset, o que mantém a velocidade do seu dataset otimizada. Você também pode reembaralhar o dataset facilmente:

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

Se você estiver usando seu conjunto de dados em múltiplas épocas, a semente efetiva para embaralhar a ordem dos shards no buffer de embaralhamento é seed + epoch. Isso facilita o reembaralhamento de um conjunto de dados 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 a iteração de um dataset no estilo mapa, basta pular os primeiros exemplos:

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

Mas, se você usar um DataLoader com um Sampler, você deve salvar o estado do seu sampler (você pode ter escrito um sampler personalizado que permite retomar).

Por outro lado, os conjuntos de dados iteráveis não oferecem acesso aleatório a um índice específico de exemplo para retomar a partir dele. Mas você pode usar IterableDataset.state_dict() e IterableDataset.load_state_dict() para retomar a partir de um checkpoint, da mesma forma que você pode fazer para modelos e otimizadores:

>>> 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 baixo dos panos, o dataset iterável mantém o controle do shard atual sendo lido e do índice do exemplo no shard atual e armazena essas informações no state_dict.

Para retomar de um checkpoint, o dataset ignora todos os shards que foram lidos anteriormente para reiniciar a partir do shard atual. Em seguida, ele lê o shard e ignora exemplos até alcançar o exemplo exato do checkpoint.

Portanto, reiniciar um dataset é bastante rápido, pois ele não relerá os shards que já foram iterados. Ainda assim, retomar um dataset geralmente não é instantâneo, pois ele precisa reiniciar a leitura a partir do início do shard atual e pular exemplos até alcançar a localização do checkpoint.

Isso pode ser usado com o StatefulDataLoader do torchdata, veja streaming com um PyTorch DataLoader.

Alterne de estilo de mapa para iterável

Se você deseja aproveitar o comportamento “lazy” de um IterableDataset ou suas vantagens de velocidade, pode converter seu Dataset de estilo map em um IterableDataset:

my_iterable_dataset = my_dataset.to_iterable_dataset()

Se você deseja embaralhar seu conjunto de dados ou usá-lo com um PyTorch DataLoader, recomendamos gerar um IterableDataset fragmentado:

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