Datasets 文档
与 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 内部重新加载数据集不会占用您的 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)
在这种情况下,每个 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()。
分布式
要跨训练节点拆分数据集,您可以使用 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 数据集和 iterable 数据集。数据集在大小为 world_size
的节点池中,为 rank rank
的节点拆分。
对于 map-style 数据集
每个节点都被分配一个数据块,例如,rank 0 被赋予数据集的第一个块。
对于 iterable 数据集
如果数据集的分片数量是 world_size
的因子(即如果 dataset.num_shards % world_size == 0
),则分片在节点之间均匀分配,这是最优化的。否则,每个节点保留 world_size
中的 1 个示例,跳过其他示例。
如果希望每个节点使用多个 worker 加载数据,这也可以与 torch.utils.data.DataLoader
结合使用。