Datasets documentation
PyTorch와 함께 사용
PyTorch와 함께 사용
이 문서는 datasets를 PyTorch와 함께 사용하는 방법에 대한 간단한 소개입니다. 특히 데이터셋에서 torch.Tensor 객체를 어떻게 가져오는지, 그리고 PyTorch DataLoader와 Hugging Face Dataset를 최상의 성능으로 사용하는 방법에 중점을 둡니다.
데이터셋 형식
기본적으로 datasets는 정수, 실수, 문자열, 리스트 등 일반적인 Python 객체를 반환합니다.
대신 PyTorch 텐서를 얻으려면 Dataset.with_format()을 사용하여 데이터셋의 형식을 pytorch로 설정할 수 있습니다:
>>> from datasets import Dataset
>>> data = [[1, 2],[3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([1, 2])}
>>> ds[:2]
{'data': tensor([[1, 2],
[3, 4]])}Dataset 객체는 Arrow 테이블의 래퍼로, 데이터셋의 배열에서 PyTorch 텐서로의 빠른 제로-복사(Zero-copy) 읽기를 허용합니다.
GPU에서 데이터를 텐서로 로드하려면 device 인수를 지정하세요:
>>> import torch
>>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
>>> ds = ds.with_format("torch", device=device)
>>> ds[0]
{'data': tensor([1, 2], device='cuda:0')}N차원 배열
데이터셋이 N차원 배열로 구성되어 있다면, 기본적으로 shape가 고정되어 있으면 동일한 tensor로 간주된다는 것을 확인할 수 있습니다:
>>> from datasets import Dataset
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]] # fixed shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([[1, 2],
[3, 4]])}>>> from datasets import Dataset
>>> data = [[[1, 2],[3]],[[4, 5, 6],[7, 8]]] # varying shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': [tensor([1, 2]), tensor([3])]}그러나 이러한 로직은 느린 shape 비교와 데이터 복사를 요구하는 경우가 많습니다. 이를 피하려면 Array feature type을 명시적으로 사용하고 텐서의 shape를 지정해야 합니다:
>>> from datasets import Dataset, Features, Array2D
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]
>>> features = Features({"data": Array2D(shape=(2, 2), dtype='int32')})
>>> ds = Dataset.from_dict({"data": data}, features=features)
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([[1, 2],
[3, 4]])}
>>> ds[:2]
{'data': tensor([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]])}기타 기능 유형
ClassLabel 데이터가 텐서로 올바르게 변환됩니다:
>>> from datasets import Dataset, Features, ClassLabel
>>> labels = [0, 0, 1]
>>> features = Features({"label": ClassLabel(names=["negative", "positive"])})
>>> ds = Dataset.from_dict({"label": labels}, features=features)
>>> ds = ds.with_format("torch")
>>> ds[:3]
{'label': tensor([0, 0, 1])}PyTorch는 숫자만 지원하므로 문자열 및 바이너리 객체는 변경되지 않습니다.
Image 기능 유형을 사용하려면
vision추가 항목을pip install datasets[vision]으로 설치해야 합니다.
>>> from datasets import Dataset, Features, Audio, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features)
>>> ds = ds.with_format("torch")
>>> ds[0]["image"].shape
torch.Size([512, 512, 4])
>>> ds[0]
{'image': tensor([[[255, 215, 106, 255],
[255, 215, 106, 255],
...,
[255, 255, 255, 255],
[255, 255, 255, 255]]], dtype=torch.uint8)}
>>> ds[:2]["image"].shape
torch.Size([2, 512, 512, 4])
>>> ds[:2]
{'image': tensor([[[[255, 215, 106, 255],
[255, 215, 106, 255],
...,
[255, 255, 255, 255],
[255, 255, 255, 255]]]], dtype=torch.uint8)}Audio 기능 유형을 사용하려면
audio추가 항목을pip install datasets[audio]로 설치해야 합니다.
>>> from datasets import Dataset, Features, Audio, Image
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features)
>>> ds = ds.with_format("torch")
>>> ds[0]["audio"]["array"]
tensor([ 6.1035e-05, 1.5259e-05, 1.6785e-04, ..., -1.5259e-05,
-1.5259e-05, 1.5259e-05])
>>> ds[0]["audio"]["sampling_rate"]
tensor(44100)데이터 로딩
torch.utils.data.Dataset 객체와 마찬가지로, Dataset는 PyTorch의 DataLoader에 직접 전달할 수 있습니다:
>>> import numpy as np
>>> from datasets import Dataset
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(16)
>>> label = np.random.randint(0, 2, size=16)
>>> ds = Dataset.from_dict({"data": data, "label": label}).with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=4)
>>> for batch in dataloader:
... print(batch)
{'data': tensor([0.0047, 0.4979, 0.6726, 0.8105]), 'label': tensor([0, 1, 0, 1])}
{'data': tensor([0.4832, 0.2723, 0.4259, 0.2224]), 'label': tensor([0, 0, 0, 0])}
{'data': tensor([0.5837, 0.3444, 0.4658, 0.6417]), 'label': tensor([0, 1, 0, 0])}
{'data': tensor([0.7022, 0.1225, 0.7228, 0.8259]), 'label': tensor([1, 1, 1, 1])}데이터 로딩 최적화
데이터 로딩 속도를 높여 시간을 절약할 수 있는 여러 가지 방법이 있습니다. 특히 대규모 데이터셋을 다룰 때 유용합니다. PyTorch는 병렬화된 데이터 로딩, 개별 인덱스 대신 인덱스 배치(batch)를 가져오는 방식, 그리고 디스크에 다운로드하지 않고 데이터셋을 반복(iterate)하기 위한 스트리밍을 제공합니다.
여러 개의 Workers 사용
PyTorch DataLoader의 num_workers 인수를 사용하여 데이터 로딩을 병렬화하면 더 높은 처리량을 얻을 수 있습니다.
내부적으로 DataLoader는 num_workers개의 프로세스를 시작합니다. 각 프로세스는 DataLoader에 전달된 데이터셋을 다시 로드하여 예시를 조회하는 데 사용됩니다. 워커 내에서 데이터셋을 다시 로드해도 RAM이 가득 차지 않습니다. 이는 단순히 디스크에서 데이터셋을 다시 메모리 맵하기 때문입니다.
>>> import numpy as np
>>> from datasets import Dataset, load_from_disk
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).save_to_disk("my_dataset")
>>> ds = load_from_disk("my_dataset").with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=32, num_workers=4)데이터 스트리밍
IterableDataset로 로드하여 데이터셋을 스트리밍합니다. 이를 통해 원격 데이터셋을 디스크에 다운로드하지 않고 점진적으로 반복하거나, 로컬 데이터 파일을 반복할 수 있습니다. 일반 데이터셋과 반복 가능 데이터셋 선택 가이드에서 사용 사례에 가장 적합한 데이터셋 유형에 대해 자세히 알아보세요.
datasets의 반복 가능한 데이터셋은 torch.utils.data.IterableDataset을 상속하므로 torch.utils.data.DataLoader에 전달할 수 있습니다:
>>> import numpy as np
>>> from datasets import Dataset, load_dataset
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).push_to_hub("<username>/my_dataset") # Upload to the Hugging Face Hub
>>> my_iterable_dataset = load_dataset("<username>/my_dataset", streaming=True, split="train")
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32)데이터셋이 여러 샤드로 나뉘어 있는 경우(즉, 데이터셋이 여러 데이터 파일로 구성된 경우) num_workers를 사용하여 병렬로 스트리밍할 수 있습니다:
>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> my_iterable_dataset.num_shards
39
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32, num_workers=4)이 경우 각 워커는 스트리밍할 샤드 목록의 부분 집합을 할당받습니다.
체크포인트 및 재개
학습 중간에 체크포인트를 저장하고 재개할 수 있는 DataLoader가 필요하면 torchdata의 StatefulDataLoader를 사용할 수 있습니다:
>>> from torchdata.stateful_dataloader import StatefulDataLoader
>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> dataloader = StatefulDataLoader(my_iterable_dataset, batch_size=32, num_workers=4)
>>> # save in the middle of training
>>> state_dict = dataloader.state_dict()
>>> # and resume later
>>> dataloader.load_state_dict(state_dict)이것은 IterableDataset.state_dict()와 IterableDataset.load_state_dict() 덕분에 가능합니다.
분산
To split your dataset across your training nodes, you can use datasets.distributed.split_dataset_by_node():
import os
from datasets.distributed import split_dataset_by_node
ds = split_dataset_by_node(ds, rank=int(os.environ["RANK"]), world_size=int(os.environ["WORLD_SIZE"]))이 방법은 map-style dataset과 iterable dataset 모두에 적용됩니다.
dataset은 world_size 크기의 노드 풀에서 rank가 rank인 노드를 위해 분할됩니다.
map 스타일 데이터셋의 경우:
각 노드는 데이터의 청크를 할당받습니다. 예를 들어, rank 0은 데이터셋의 첫 번째 청크를 받습니다.
반복 가능한 데이터셋의 경우:
strategy 분할은 "auto", "shards", 또는 "examples"일 수 있습니다:
"shards": 샤드는 노드에 고르게 할당됩니다"examples": 각 노드는world_size개 중 1개의 예제를 유지하고, 나머지 예제는 건너뜁니다"auto"(기본값): 데이터셋의 샤드 수가world_size의 약수인 경우(즉,dataset.num_shards % world_size == 0인 경우) 가장 최적화된"shards"를 사용합니다. 그렇지 않으면"examples"를 사용합니다.
각 노드가 여러 워커를 사용하여 데이터를 로드하려면 torch.utils.data.DataLoader와 함께 사용할 수도 있습니다.
분산 환경에서 반복 가능(iterable) 데이터셋을 섞을 경우, 모든 노드에서 동일한 섞인 샤드 목록을 사용하여 각 노드가 건너뛰어야 할 샤드를 파악할 수 있도록 IterableDataset.shuffle()에서 고정된
seed를 설정해야 합니다.
하이브리드 샤드 데이터 병렬(HSDP)
torch.distributed.device_mesh.DeviceMesh는 학습 토폴로지를 랭크의 그리드로 설명하며, 각 병렬화 전략마다 하나의 차원을 가집니다. 데이터 로딩에는 데이터 병렬 차원만 중요하며, 하이브리드 샤딩을 사용하면 이를 놓치기 쉽습니다:
"fsdp"(일반적으로 노드 내): 각 랭크가 모델 상태(파라미터, 그래디언트, 옵티마이저 상태)의 샤드만 유지하고, 해당 모듈에 대한 계산을 수행하기 직전에 해당 모듈의 전체 파라미터를 일시적으로 올-거더하는 FSDP 그룹입니다. 학습은 데이터 병렬로 유지됩니다. 각 랭크가 자신의 배치에 대해 전체 포워드 및 백워드 패스를 실행하므로, DDP와 마찬가지로 이 차원에서는 랭크마다 서로 다른 데이터가 필요합니다."dp_replicate"(일반적으로 노드 간): 해당 샤드된 모델의 중복된 복사본으로, 리플리카 간에 추가적인 all-reduce를 통해 동기화됩니다. 리플리카는 서로 다른 데이터를 처리합니다."tp"(intra-node): 각 레이어의 가중치를 분할하여 하나의 배치 계산을 함께 수행하는 tensor parallel rank입니다. data-parallel이 아니므로 동일한 예시를 읽어야 합니다.
데이터 병렬 도(sharding(FSDP) 차원과 복제 차원의 곱)는 메시에 "tp"와 같은 차원이 존재하는 순간 WORLD_SIZE보다 작아집니다. datasets.distributed.split_dataset_by_node()에 이 곱과 두 차원에 걸쳐 평탄화된 랭크를 전달하세요:
from torch.distributed.device_mesh import init_device_mesh
from datasets.distributed import split_dataset_by_node
# Hybrid sharding spreads data parallelism over two mesh dimensions, but both
# need different data per rank, so `dp_size` is their product.
mesh = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("dp_replicate", "fsdp", "tp"))
dp_replicate_mesh, fsdp_mesh = mesh["dp_replicate"], mesh["fsdp"]
dp_size = dp_replicate_mesh.size() * fsdp_mesh.size()
dp_rank = dp_replicate_mesh.get_local_rank() * shard.size() + fsdp_mesh.get_local_rank()
ds = split_dataset_by_node(ds, rank=dp_rank, world_size=dp_size)여기서 dp_size는 2 * 2 = 4입니다. 각 데이터 병렬 랭크는 데이터셋의 4분의 1을 읽고, 각 "tp" 쌍의 두 랭크는 동일한 예시를 읽습니다.