Datasets documentation
与 PyTorch 配合使用
与 PyTorch 配合使用
本文档简要介绍了如何使用 datasets 与 PyTorch 配合，重点说明如何从数据集中获取 torch.Tensor 对象，以及如何结合 PyTorch DataLoader 和 Hugging Face Dataset 以获得最佳性能。
数据集格式
默认情况下，数据集返回的是普通的 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 张量。
要在 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 维数组组成，您会发现默认情况下，如果形状固定，它们会被视为相同的张量：
>>> 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])]}然而，这种逻辑通常需要缓慢的形状比较和数据复制。
为避免这种情况，您必须显式使用 Array 特征类型，并指定张量的形状：
>>> 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 提供了并行化数据加载功能，可以批量获取索引而非逐个获取，并支持流式处理，以便在无需将数据集下载到磁盘的情况下对其进行迭代。
使用多个 Workers
你可以使用 PyTorch DataLoader 的 num_workers 参数来并行化数据加载，从而获得更高的吞吐量。
在底层，DataLoader 会启动 num_workers 个进程。
每个进程都会重新加载传递给 DataLoader 的数据集，并用于查询样本。
在 worker 内部重新加载数据集不会占满你的内存，因为它只是再次从磁盘对数据集进行内存映射。
>>> 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)在这种情况下，每个 worker 都会被分配一个分片列表的子集，用于从中流式读取数据。
检查点与恢复
如果你需要一个可以在训练中途进行检查点和恢复的 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 风格数据集和可迭代数据集。
数据集将按照节点池大小 world_size 中 rank 为 rank 的节点进行划分。
对于映射式数据集：
每个节点会被分配一块数据，例如 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 结合使用。
如果在分布式环境中对可迭代数据集进行洗牌，请确保在 IterableDataset.shuffle() 中设置固定的
seed，以便在每个节点上使用相同的分片洗牌列表，从而确定该节点应跳过哪些分片。
混合分片数据并行（HSDP）
torch.distributed.device_mesh.DeviceMesh 以 rank 网格的形式描述您的训练拓扑结构，每个并行策略对应一个维度。对于数据加载而言，只有数据并行维度是重要的，而混合分片策略很容易让人忽略这些维度：
"fsdp"（通常用于节点内）：FSDP 组中，每个 rank 仅保留模型状态（参数、梯度和优化器状态）的一个分片，并在对某个模块进行计算前临时 all-gather 该模块的完整参数。训练保持数据并行：每个 rank 在自己的 batch 上执行完整的前向和反向传播，因此该维度需要每个 rank 拥有不同的数据，与 DDP 相同。"dp_replicate"（通常用于节点间）：该分片模型的重复副本，通过副本间的额外 all-reduce 操作保持同步。各副本还会处理不同的数据。"tp"（节点内）：张量并行秩，将每一层的权重拆分以共同计算一个批次。它们不是数据并行的，必须读取相同的样本。
数据并行度是分片（FSDP）维度和复制维度的乘积，只要网格中存在诸如 "tp" 这样的维度，该乘积就会小于 WORLD_SIZE。请将此乘积以及在这两个维度上展平后的 rank 传递给 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：每个数据并行 rank 读取数据集的四分之一，而每个 "tp" 对中的两个 rank 读取相同的样本。