WebDataset
WebDataset 是一个用于编写大型数据集 I/O 管道的库。其顺序 I/O 和分片功能使其特别适用于将大规模数据集流式传输到 DataLoader。
WebDataset 格式
WebDataset 文件是一个 TAR 归档文件,其中包含一系列数据文件。所有具有相同前缀的连续数据文件都被视为属于同一个示例(例如,图像/音频文件及其标签或元数据)
标签和元数据可以存储在 .json
文件中,也可以存储在 .txt
文件中(用于标题、描述),或者存储在 .cls
文件中(用于类索引)。
一个大型 WebDataset 由许多称为分片的文件组成,每个分片都是一个 TAR 归档文件。每个分片通常约为 1GB,但完整的数据集可以达到多个 TB!
流式传输
流式传输 TAR 归档文件速度很快,因为它读取连续的数据块。它可能比逐个读取单独的数据文件快几个数量级。
WebDataset 流式传输在从磁盘和云存储读取时都提供了高速性能,这使其成为馈送到 DataLoader 的理想格式。
例如,以下是如何直接从 Hugging Face 流式传输 timm/imagenet-12k-wds 数据集。
首先,您需要 使用您的 Hugging Face 帐户登录,例如使用
huggingface-cli login
然后,您可以使用 WebDataset 流式传输数据集。
>>> import webdataset as wds
>>> from huggingface_hub import get_token
>>> from torch.utils.data import DataLoader
>>> hf_token = get_token()
>>> url = "https://huggingface.co/datasets/timm/imagenet-12k-wds/resolve/main/imagenet12k-train-{{0000..1023}}.tar"
>>> url = f"pipe:curl -s -L {url} -H 'Authorization:Bearer {hf_token}'"
>>> dataset = wds.WebDataset(url).decode()
>>> dataloader = DataLoader(dataset, batch_size=64, num_workers=4)
洗牌
通常,WebDataset 格式中的数据集已经过洗牌并准备馈送到 DataLoader。但是您仍然可以使用 WebDataset 的近似洗牌重新洗牌数据。
除了洗牌分片列表之外,WebDataset 还使用缓冲区在没有任何速度损失的情况下洗牌数据集。
要洗牌分片文件列表并从洗牌缓冲区中随机采样
>>> buffer_size = 1000
>>> dataset = (
... wds.WebDataset(url, shardshuffle=True)
... .shuffle(buffer_size)
... .decode()
... )