Dataset과 IterableDataset의 차이

데이터셋 객체는 Dataset과 IterableDataset 두 가지 유형이 있습니다. 사용하거나 생성할 데이터셋 유형은 데이터셋의 크기에 따라 달라집니다. 일반적으로 IterableDataset는 지연 평가(lazy) 동작과 속도상의 이점으로 인해 대규모 데이터셋(수백 GB 규모!)에 이상적이며, Dataset은 그 외 모든 경우에 적합합니다. 이 페이지에서는 Dataset과 IterableDataset의 차이를 비교하여 사용자에게 적합한 데이터셋 객체를 선택하는 데 도움을 드립니다.

다운로드 및 스트리밍

일반적인 Dataset이 있는 경우, my_dataset[0]을 사용하여 액세스할 수 있습니다. 이를 통해 행에 대한 랜덤 액세스가 가능합니다. 이러한 데이터셋은 “map-style” 데이터셋이라고도 합니다. 예를 들어, 다음과 같이 ImageNet-1k를 다운로드하고 임의의 행에 액세스할 수 있습니다:

from datasets import load_dataset

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

다만, 전체 데이터셋을 디스크나 메모리에 저장해야 하므로 디스크보다 큰 데이터셋에 접근하는 것이 불가능하다는 단점이 있습니다. 대용량 데이터셋에서 이러한 불편함이 발생할 수 있어, IterableDataset이라는 다른 유형의 데이터셋이 존재합니다. IterableDataset을 사용하면 for 루프를 통해 데이터셋을 반복하면서 데이터를 점진적으로 로드할 수 있습니다. 이 방식은 메모리에 예시의 일부만 로드하고 디스크에 아무것도 기록하지 않습니다.

예를 들어, ImageNet-1k 데이터셋을 디스크에 다운로드하지 않고 스트리밍할 수 있습니다:

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은 디스크에 파일을 쓰지 않고 온라인 데이터를 읽을 수 있습니다. 예를 들어, C4나 LAION-2B처럼 각 샤드가 수백 기가바이트인 여러 샤드로 구성된 데이터셋을 스트리밍할 수 있습니다. 데이터셋 스트리밍 방법에 대한 자세한 내용은 Dataset Streaming Guide에서 확인하세요.

그러나 이것이 유일한 차이는 아닙니다. IterableDataset의 “lazy” 동작은 데이터셋 생성 및 처리 과정에서도 동일하게 나타납니다.

맵 스타일 데이터셋과 반복 가능 데이터셋 만들기

리스트나 딕셔너리를 사용하여 Dataset을 생성할 수 있으며, 데이터는 완전히 Arrow로 변환되어 어떤 행이든 쉽게 접근할 수 있습니다:

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

반면에 IterableDataset을 생성하려면 데이터를 “lazy”하게 로드하는 방법을 제공해야 합니다. Python에서는 일반적으로 생성자 함수를 사용합니다. 이러한 함수는 yield를 사용하여 한 번에 하나의 예시를 반환하므로, 일반적인 Dataset처럼 슬라이싱으로 행에 접근할 수 없습니다:

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

로컬 파일을 완전히 또는 점진적으로 로드하기

It is possible to convert local or remote data files to an Arrow Dataset using load_dataset():

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

그러나 이는 CSV에서 Arrow 형식으로의 변환 단계를 필요로 하며, 데이터셋이 크다면 시간과 디스크 공간이 많이 소요됩니다.

디스크 공간을 절약하고 변환 단계를 건너뛰려면, 로컬 파일에서 직접 스트리밍하여 IterableDataset을 정의할 수 있습니다. 이렇게 하면 데이터셋을 반복하는 동안 로컬 파일에서 데이터를 점진적으로 읽게 됩니다:

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

CSV, JSONL, Parquet와 같은 다양한 파일 형식 및 이미지와 오디오 파일을 지원합니다. 테이블, 텍스트, 비전, 오디오 데이터셋을 로드하는 방법에 대한 자세한 내용은 각 가이드를 참고하세요.

Eager 데이터 처리와 lazy 데이터 처리

Dataset 객체를 Dataset.map()으로 처리하면 전체 데이터셋이 즉시 처리되어 반환됩니다. 예를 들어 pandas의 동작 방식과 유사합니다.

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

반면, IterableDataset의 “lazy” 특성으로 인해 IterableDataset.map()을 호출해도 map 함수가 전체 데이터셋에 적용되지 않습니다. 대신, map 함수가 실시간으로 적용됩니다.

이 때문에 여러 처리 단계를 체인할 수 있으며, 데이터셋을 반복하기 시작하면 모든 단계가 동시에 실행됩니다:

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

정확하고 빠른 근사 셔플링

Dataset를 Dataset.shuffle()로 섞으면 데이터셋을 정확히 섞는 작업을 수행합니다. 이는 인덱스 목록 [0, 1, 2, ... len(my_dataset) - 1]을 가져와 이 목록을 섞는 방식으로 작동합니다. 그 후 my_dataset[0]에 접근하면 섞인 인덱스 매핑의 첫 번째 요소로 정의된 행과 인덱스가 반환됩니다:

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

IterableDataset의 경우 행에 대한 랜덤 접근이 불가능하므로, 셔플된 인덱스 목록을 사용해 임의의 위치에서 행에 접근할 수 없습니다. 이로 인해 정확한 셔플링을 사용할 수 없습니다. 대신 IterableDataset.shuffle()에서는 빠른 근사 셔플링을 사용합니다. 셔플 버퍼를 사용해 데이터셋에서 반복적으로 랜덤한 예시를 샘플링합니다. 데이터셋이 여전히 반복적으로 읽히므로, 뛰어난 속도 성능을 제공합니다:

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

하지만 shuffle buffer만으로는 머신러닝 모델 학습에 만족스러운 shuffling을 제공하기 어렵습니다. 따라서 IterableDataset.shuffle()는 데이터셋이 여러 파일이나 소스로 구성되어 있는 경우 데이터셋 샤드도 shuffling합니다. 또한, 버퍼를 채울 때 한 번에 최대 max_buffer_input_shards개의 샤드(기본값은 10)를 사용하므로 shuffling의 품질이 크게 향상됩니다:

# 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

데이터셋의 파일 수가 적다면, reshard()를 셔플링 전에 호출하여 셔플링 품질을 향상시킬 수 있습니다. 예를 들어, Parquet 데이터셋의 경우 파일 하나당 샤드 하나를 사용하는 대신 row group을 기준으로 리샤딩하여 셔플링할 수 있는 샤드 수를 크게 늘릴 수 있습니다.

속도 차이

일반적인 Dataset 객체는 Arrow를 기반으로 하며, 이를 통해 행에 대한 빠른 랜덤 액세스를 제공합니다. 메모리 매핑과 Arrow가 인메모리 형식이라는 특성 덕분에 디스크에서 데이터를 읽을 때 비용이 높은 시스템 호출과 역직렬화가 발생하지 않습니다. 연속된 Arrow 레코드 배치에 대해 반복함으로써 for 루프를 사용한 반복 시 더 빠른 데이터 로딩을 제공합니다.

하지만 Dataset에 인덱스 매핑이 있는 경우(예: Dataset.shuffle()를 통해) 속도가 10배까지 느려질 수 있습니다. 이것은 인덱스 매핑을 사용하여 읽을 행 인덱스를 가져오는 추가 단계가 필요하기 때문이며, 가장 중요한 것은 더 이상 연속된 데이터 청크를 읽지 않기 때문입니다. 속도를 회복하려면 인덱스 매핑을 제거하는 Dataset.flatten_indices()를 사용하여 디스크에 전체 데이터셋을 다시 작성해야 합니다. 다만 데이터셋의 크기에 따라 상당한 시간이 소요될 수 있습니다:

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

이 경우 IterableDataset으로 전환하고, 그 빠른 근사 셔플링 메서드인 IterableDataset.shuffle()을 활용하는 것을 권장합니다. 이 메서드는 샤드 순서만 셔플링하고 데이터셋에 셔플 버퍼를 추가하므로, 데이터셋의 속도를 최적 상태로 유지합니다. 데이터셋을 쉽게 다시 셔플링할 수도 있습니다:

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

데이터셋을 여러 epoch에 걸쳐 사용하는 경우, shuffle buffer에서 shard 순서를 섞기 위한 유효한 시드는 seed + epoch입니다. 이를 통해 epoch 간에 데이터셋을 다시 섞는 것이 쉬워집니다:

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

map 스타일 데이터셋의 반복을 다시 시작하려면 첫 번째 예시를 건너뛰면 됩니다:

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

하지만 DataLoader에 Sampler를 사용하는 경우, 대신 샘플러의 상태를 저장해야 합니다(재개를 허용하는 커스텀 샘플러를 작성했을 수 있습니다).

반면, iterable dataset은 특정 예시 인덱스에 대한 랜덤 액세스를 제공하지 않으므로 해당 위치에서 재개할 수 없습니다. 대신 모델과 옵티마이저에서 할 수 있는 것과 유사하게 IterableDataset.state_dict()와 IterableDataset.load_state_dict()를 사용하여 체크포인트에서 재개할 수 있습니다:

>>> 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)

내부적으로, iterable dataset은 현재 읽는 중인 shard와 해당 shard 내의 예시 인덱스를 추적하며, 이 정보를 state_dict에 저장합니다.

체크포인트에서 재개하려면, 데이터셋은 이전에 읽은 모든 샤드를 건너뛰고 현재 샤드에서 다시 시작합니다. 그 다음 샤드를 읽고, 체크포인트의 정확한 예제에 도달할 때까지 예제를 건너뜁니다.

따라서 데이터셋을 다시 시작하는 것은 이미 반복된 샤드를 다시 읽지 않으므로 상당히 빠릅니다. 여전히 데이터셋을 재개하는 것은 일반적으로 즉각적이지 않으며, 현재 샤드의 처음부터 다시 읽기 시작하여 체크포인트 위치에 도달할 때까지 예제를 건너뛰어야 하기 때문입니다.

torchdata의 StatefulDataLoader와 함께 사용할 수 있으며, PyTorch DataLoader를 사용한 스트리밍을 참고하세요.

map 스타일에서 iterable로 전환

IterableDataset의 “lazy” 동작이나 속도 이점을 활용하려면, map-style Dataset을 IterableDataset으로 전환할 수 있습니다:

my_iterable_dataset = my_dataset.to_iterable_dataset()

데이터셋을 섞거나 PyTorch DataLoader와 함께 사용하려면, 샤드된 IterableDataset을 생성하는 것을 권장합니다:

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