데이터셋 스트리밍

데이터셋 스트리밍 (Stream)

데이터셋 스트리밍을 사용하면 데이터셋을 다운로드하지 않고 작업할 수 있어요. 데이터는 데이터셋을 순회하는 동안 즉석에서 스트리밍되죠. 이 방법은 특히 다음 상황에서 유용해요:

엄청나게 큰 데이터셋의 다운로드를 기다리기 싫을 때
데이터셋 크기가 컴퓨터 사용 가능한 디스크 공간을 초과할 때
데이터셋의 몇 개 샘플만 빠르게 탐색하고 싶을 때

예를 들어 HuggingFaceFW/fineweb 데이터셋의 영어 split은 45테라바이트지만, 스트리밍으로는 즉시 사용할 수 있어요. 아래처럼 load_dataset()에서 streaming=True로 설정하면 데이터셋을 스트리밍할 수 있습니다.

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> print(next(iter(dataset)))
{'text': 'How AP reported in all formats from tornado-stricken regionsMarch 8, 2012\nWhen the first serious bout of tornadoes of 2012 blew through middle America in the middle of the night, they touched down in places hours from any AP bureau...', ...,
 'language_score': 0.9721424579620361, 'token_count': 717}

데이터셋 스트리밍은 로컬 파일로 이루어진 데이터셋도 변환 없이 작업할 수 있게 해 줘요. 이 경우 데이터는 데이터셋을 순회하는 동안 로컬 파일에서 스트리밍됩니다. 특히 다음 상황에서 유용해요:

엄청나게 큰 로컬 데이터셋이 Arrow로 변환되기를 기다리기 싫을 때
변환된 파일 크기가 컴퓨터의 사용 가능한 디스크 공간을 초과할 때
데이터셋의 몇 개 샘플만 빠르게 탐색하고 싶을 때
특정 컬럼만 불러오거나 Parquet 데이터셋을 효율적으로 필터링하고 싶을 때

예를 들어 oscar-corpus/OSCAR-2201 같은 수백 개의 압축 JSONL 파일로 이루어진 로컬 데이터셋을 스트리밍해서 즉시 사용할 수 있어요.

>>> from datasets import load_dataset
>>> data_files = {'train': 'path/to/OSCAR-2201/compressed/en_meta/*.jsonl.gz'}
>>> dataset = load_dataset('json', data_files=data_files, split='train', streaming=True)
>>> print(next(iter(dataset)))
{'id': 0, 'text': 'Founded in 2015, Golden Bees is a leading programmatic recruitment platform dedicated to employers, HR agencies and job boards. The company has developed unique HR-custom technologies and predictive algorithms to identify and attract the best candidates for a job opportunity.', ...

Parquet은 칼럼형 형식이라 컬럼의 일부만 스트리밍·로드하고 원치 않는 컬럼은 무시할 수 있어요. Parquet은 또한 (파일·행 그룹 수준의) 컬럼 통계 같은 메타데이터를 저장해 효율적인 필터링을 가능하게 해 줘요. datasets.packaged_modules.parquet.ParquetConfig의 columns와 filters 인자로 Parquet 데이터셋을 스트리밍하고, 컬럼을 선택하고, 필터를 적용할 수 있어요.

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True, columns=["url", "date"])
>>> print(next(iter(dataset)))
{'url': 'http://%[email protected]/Content/Press-Release/2012/How-AP-reported-in-all-formats-from-tornado-stricken-regions', 'date': '2013-05-18T05:48:54Z'}
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True, filters=[("language_score", ">=", 0.99)])
>>> print(next(iter(dataset)))
{'text': 'Everyone wishes for something. And lots of people believe they know how to make their wishes come true with magical thinking.\nWhat is it? "Magical thinking is a belief in forms of causation, with no known physical basis," said Professor Emily Pronin of Princeton...', ...,
 'language_score': 0.9900368452072144, 'token_count': 716}

스트리밍 모드로 데이터셋을 불러오면 (기존 Dataset 객체 대신) IterableDataset이라는 새로운 데이터셋 유형 인스턴스가 만들어져요. 이 특별한 유형의 데이터셋은 아래 보여주는 자체 처리 메서드 세트를 가져요.

IterableDataset은 모델 훈련 같은 반복(iterative) 작업에 유용해요. 하지만 IterableDataset을 예시에 랜덤 접근이 필요한 작업에는 쓰면 안 돼요. for 루프로 전부 순회해야 하기 때문이에요. 이터러블 데이터셋의 마지막 예시를 얻으려면 앞선 모든 예시를 순회해야 합니다. 자세한 내용은 Dataset vs. IterableDataset 가이드에서 볼 수 있어요.

컬럼 인덱싱

특정 컬럼의 값들을 순회하는 것이 편리할 때가 있어요. 다행히 IterableDataset은 컬럼 인덱싱을 지원해요.

>>> from datasets import load_dataset
>>> dataset = load_dataset("allenai/c4", "en", streaming=True, split="train")
>>> print(next(iter(dataset["text"])))
Beginners BBQ Class Taking Place in Missoula!...

Dataset에서 변환

이미 Dataset 객체가 있다면 to_iterable_dataset() 함수로 IterableDataset로 변환할 수 있어요. 이 방식은 load_dataset()에 streaming=True 인자를 설정하는 것보다 실제로 더 빨라요. 데이터가 로컬 파일에서 스트리밍되기 때문이에요.

>>> from datasets import load_dataset

# 더 빠름 🐇
>>> dataset = load_dataset("ethz/food101")
>>> iterable_dataset = dataset.to_iterable_dataset()

# 더 느림 🐢
>>> iterable_dataset = load_dataset("ethz/food101", streaming=True)

to_iterable_dataset() 함수는 IterableDataset이 인스턴스화될 때 샤딩을 지원해요. 큰 데이터셋을 다룰 때 유용하며, 데이터셋을 셔플하거나 PyTorch DataLoader로 빠른 병렬 로딩을 켜고 싶을 때 씁니다.

>>> import torch
>>> from datasets import load_dataset

>>> dataset = load_dataset("ethz/food101")
>>> iterable_dataset = dataset.to_iterable_dataset(num_shards=64) # 데이터셋 샤딩
>>> iterable_dataset = iterable_dataset.shuffle(buffer_size=10_000)  # 순회를 시작할 때 샤드 순서를 셔플하고 셔플 버퍼 사용
dataloader = torch.utils.data.DataLoader(iterable_dataset, num_workers=4)  # 순회 시작 시 셔플된 샤드 리스트에서 각 워커에 64 / 4 = 16개 샤드 할당

셔플

일반 Dataset 객체처럼 IterableDataset도 IterableDataset.shuffle()로 셔플할 수 있어요.

buffer_size 인자는 예시를 무작위로 샘플링할 버퍼의 크기를 제어해요. 데이터셋에 백만 개의 예시가 있고 buffer_size를 만으로 설정했다고 해 볼게요. IterableDataset.shuffle()은 버퍼의 첫 만 개 예시에서 무작위로 예시를 선택할 거예요. 버퍼에서 선택된 예시는 새 예시로 교체됩니다. 기본 버퍼 크기는 1,000이에요.

>>> from datasets import load_dataset
>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> shuffled_dataset = dataset.shuffle(seed=42, buffer_size=10_000)

IterableDataset.shuffle()은 데이터셋이 여러 파일로 샤딩되어 있다면 샤드의 순서도 셔플해요.

리셔플 (Reshuffle)

때로는 에폭마다 데이터셋을 다시 셔플하고 싶을 때가 있어요. 그러려면 에폭마다 다른 시드를 설정해야 해요. 에폭 사이에 IterableDataset.set_epoch()를 사용해 현재 몇 번째 에폭인지 데이터셋에 알려주세요.

시드는 실질적으로 초기 시드 + 현재 에폭이 됩니다.

>>> for epoch in range(epochs):
...     shuffled_dataset.set_epoch(epoch)
...     for example in shuffled_dataset:
...         ...

데이터셋 분할

데이터셋을 두 가지 방법 중 하나로 나눌 수 있어요.

IterableDataset.take()는 데이터셋의 처음 n개 예시를 반환해요.

>>> dataset = load_dataset('HuggingFaceFW/fineweb', split='train', streaming=True)
>>> dataset_head = dataset.take(2)
>>> list(dataset_head)
[{'text': "How AP reported in all formats from tor...},
 {'text': 'Did you know you have two little yellow...}]

IterableDataset.skip()은 데이터셋의 처음 n개 예시를 건너뛰고 나머지 예시를 반환해요.

>>> train_dataset = shuffled_dataset.skip(1000)

take와 skip은 샤드의 순서를 고정하므로 이후에 shuffle을 호출하는 걸 막아요. 데이터셋을 나누기 전에 셔플해야 합니다.

샤드

🤗 Datasets은 아주 큰 데이터셋을 미리 정한 개수의 청크로 나누는 샤딩을 지원해요. shard()의 num_shards 파라미터로 데이터셋을 몇 개의 샤드로 나눌지 정하고, 반환할 샤드를 index 파라미터로 지정하면 됩니다.

예를 들어 amazon_polarity 데이터셋에는 4개의 샤드(이 경우 Parquet 파일 4개)가 있어요.

>>> from datasets import load_dataset
>>> dataset = load_dataset("fancyzhx/amazon_polarity", split="train", streaming=True)
>>> print(dataset)
IterableDataset({
    features: ['label', 'title', 'content'],
    num_shards: 4
})

데이터셋을 두 개의 청크로 샤딩하면 첫 번째는 2개 샤드만 가집니다.

>>> dataset.shard(num_shards=2, index=0)
IterableDataset({
    features: ['label', 'title', 'content'],
    num_shards: 2
})

데이터셋의 샤드 수를 늘리려면 IterableDataset.reshard()를 사용할 수 있어요.

체크포인트 재개

IterableDataset을 순회하는 동안 어느 샤드·어느 예시까지 진행했는지를 기록해 두는 것이 가능해요. 이 정보는 state_dict 형태로 저장하거나 불러올 수 있어요.

따라서 데이터셋을 다시 시작하는 것은 꽤 빠른데, 이미 순회한 샤드를 다시 읽지 않기 때문이에요. 그래도 데이터셋 재개는 일반적으로 즉각적이지 않아요. 현재 샤드의 처음부터 다시 읽고 체크포인트 위치에 도달할 때까지 예시를 건너뛰어야 하기 때문이에요.

이것은 torchdata의 StatefulDataLoader와 함께 사용할 수 있어요.

>>> from torchdata.stateful_dataloader import StatefulDataLoader
>>> iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> dataloader = StatefulDataLoader(iterable_dataset, batch_size=32, num_workers=4)
>>> # 체크포인트
>>> state_dict = dataloader.state_dict()  # 내부적으로 iterable_dataset.state_dict() 사용
>>> # 체크포인트부터 재개
>>> dataloader.load_state_dict(state_dict)  # 내부적으로 iterable_dataset.load_state_dict() 사용

재개는 체크포인트가 저장된 지점을 정확히 돌아주는데, .shuffle()을 사용한 경우는 예외예요. 셔플 버퍼의 예시들은 재개 시 손실되고 버퍼는 새 데이터로 다시 채워집니다.

저장 (Save)

이터러블 데이터셋이 준비되면 Parquet 형식의 Hugging Face Dataset으로 저장해서 load_dataset()으로 나중에 다시 쓸 수 있어요.

push_to_hub()에 데이터셋을 저장할 Hugging Face 데이터셋 저장소 이름을 제공해서 저장하세요. 이 방식은 데이터셋을 순회하며 데이터를 Hugging Face에 점진적으로 업로드해요.

dataset.push_to_hub("username/my_dataset")

데이터셋이 여러 샤드(dataset.num_shards > 1)로 구성되어 있다면 여러 프로세스로 병렬 업로드할 수 있어요. 특히 map()이나 filter() 단계를 적용했다면 병렬로 더 빨리 실행되므로 유용합니다.

dataset.push_to_hub("username/my_dataset", num_proc=8)

load_dataset() 함수로 데이터셋을 다시 불러와요.

from datasets import load_dataset
reloaded_dataset = load_dataset("username/my_dataset")

내보내기 (Export)

🤗 Datasets은 내보내기도 지원해서 다른 애플리케이션에서 데이터셋을 다룰 수 있어요. 아래 표는 현재 지원되는 내보내기 파일 형식이에요.

파일 형식	내보내기 메서드
CSV	IterableDataset.to_csv()
JSON	IterableDataset.to_json()
Parquet	IterableDataset.to_parquet()
SQL	IterableDataset.to_sql()
메모리 내 파이썬 객체	IterableDataset.to_pandas(), `IterableDataset.to_polars()`, 혹은 IterableDataset.to_dict()

예를 들어 데이터셋을 CSV 파일로 내보내려면 이렇게 해요.

>>> dataset.to_csv("path/of/my/dataset.csv")

큰 데이터셋이라면 샤드당 파일 하나로 저장할 수 있어요, 예:

>>> num_shards = dataset.num_shards
>>> for index in range(num_shards):
...     shard = dataset.shard(index, num_shards)
...     shard.to_parquet(f"path/of/my/dataset/data-{index:05d}.parquet")

출처: 공식문서