由 Alex Aizman、Gavin Maltby、Thomas Breuel 撰写

数据集每天都在增长，GPU 也在变得越来越快。这意味着深度学习研究人员和工程师有更多的数据集来训练和验证他们的模型。

  • 许多用于静态图像识别的研究数据集正在变得可用，包括包含 1000 万或更多图像的 OpenImages 和 Places。
  • 1000 万 YouTube 视频（YouTube 8M）在 720p 下消耗约 300TB，用于物体识别、视频分析和动作识别的研究。
  • 烟草语料库包含约 2000 万页扫描的高清页面，对 OCR 和文本分析研究非常有用。

尽管目前最常见的庞大数据集涉及图像和视频，但大数据集出现在许多其他领域，并涉及许多其他类型的数据：网页、金融交易、网络轨迹、脑部扫描等。

然而，处理大量数据集带来了一系列挑战：

  • 数据集大小：数据集通常超过节点本地磁盘存储的容量，需要分布式存储系统和高效的网络访问。
  • 文件数量：数据集通常包含数十亿个文件，具有均匀的随机访问模式，这往往会使本地和网络文件系统不堪重负。
  • 数据速率：在大数据集上训练作业通常需要许多 GPU，需要向数据集的聚合 I/O 带宽达到许多 GBytes/s；这只能由大规模并行 I/O 系统来满足。
  • 打乱和增强：在训练之前需要对训练数据进行打乱和增强。
  • 可扩展性：用户通常希望在小型数据集上开发和测试，然后快速扩展到大型数据集。

传统本地和网络文件系统，甚至对象存储服务器，都不是为这些应用而设计的。PyTorch 的 WebDataset I/O 库，以及可选的 AIStore 服务器和 Tensorcom RDMA 库，提供了一种高效、简单、基于标准的解决方案，以解决所有这些问题。该库足够简单，适合日常使用，基于成熟的开源标准，并且易于从现有的基于文件的数据库迁移过来。

使用 WebDataset 非常简单，几乎不需要任何努力，它可以让您将相同的代码从运行本地实验扩展到在集群或云中使用数百个 GPU，并且性能线性可扩展。即使在小型问题上和您的桌面上，它也可以将 I/O 速度提高十倍，并简化大型数据集的数据管理和处理。本博客的其余部分将向您介绍如何开始使用 WebDataset 以及它是如何工作的。

WebDataset 库

WebDataset 库为上述挑战提供了一个简单的解决方案。目前，它作为一个独立的库（github.com/tmbdev/webdataset）提供，但正在计划将其集成到 PyTorch 中（参见 RFC 38419）。WebDataset 的实现很小（大约 1500 行代码）且没有外部依赖。

WebDataset 不是发明新的格式，而是将大型数据集表示为 POSIX tar 存档文件的集合，这些存档文件包含原始数据文件。WebDataset 库可以直接使用这样的 tar 存档进行训练，无需解包或本地存储。

WebDataset 能够完美地从小型本地数据集扩展到 PB 级数据集，并在数百个 GPU 上进行训练，允许数据存储在本地磁盘、Web 服务器或专用文件服务器上。对于基于容器的训练，WebDataset 消除了对卷插件或节点本地存储的需求。作为额外的好处，在训练之前无需解包数据集，简化了研究数据的分发和使用。

WebDataset 实现了 PyTorch 的 IterableDataset 接口，可以像现有的基于 DataLoader 的代码一样使用。由于数据存储在存档内的文件中，现有的加载和数据增强代码通常只需要进行最小修改。

WebDataset 库是 PyTorch（以及也支持通过它们的 Python API 使用 TensorFlow、Keras 和 DALI）中处理大型数据集和分布式训练的完整解决方案。由于 POSIX tar 存档是一种标准、广泛支持的格式，因此很容易编写其他工具来操作这种格式的数据集。例如，tarp 命令是用 Go 编写的，可以打乱并处理训练数据集。

优点

对于非常大的数据集，使用分片、顺序可读的格式是必不可少的。此外，它在许多其他环境中也有好处。WebDataset 提供了一种解决方案，可以从桌面机器上的小问题扩展到集群或云中的非常大的深度学习问题。以下表格总结了在不同环境中的一些优点。

环境 WebDataset 的优势
带有 AIStore 的本地集群 AIStore 可以轻松作为 K8s 容器部署，提供线性可扩展性和接近 100% 的网络和 I/O 带宽利用率。适用于 petascale 深度学习。
云计算 WebDataset 深度学习作业可以直接针对存储在云存储桶中的数据集进行训练；无需卷插件。本地和云作业工作方式相同。适用于佩塔级学习。
带有现有分布式文件系统或对象存储的本地集群 WebDataset 的大规模顺序读取提高了现有分布式存储的性能，并消除了对专用卷插件的需求。
教育环境 Web 数据集可以存储在现有的 Web 服务器和 Web 缓存中，学生可以直接通过 URL 访问
在本地驱动器上对工作站进行训练 在数据下载的同时，作业可以开始训练。数据不需要解压缩即可用于训练。硬盘 I/O 性能比基于随机访问文件的数据集提高十倍。
所有环境 数据集以归档格式表示，并包含文件类型等元数据。数据以原生格式（JPEG、MP4 等）压缩。数据管理、ETL 风格的作业以及数据转换和 I/O 简化，并易于并行化。

在接下来的几个月中，我们将添加更多示例，提供基准测试，并展示如何在这些环境中使用 WebDataset。

高性能

在本地集群上进行高性能计算时，配套的开源 AIStore 服务器提供了全盘到 GPU 的 I/O 带宽，仅受硬件限制。这篇 Bigdata 2019 论文包含了详细的基准测试和性能测量。除了基准测试之外，NVIDIA 和微软的研究项目也使用了 WebDataset 来处理 PB 级数据集和数十亿的训练样本。

下面是使用 12 个服务器节点，每个节点有 10 个旋转硬盘的 AIStore 与 WebDataset 客户端的基准测试。

左轴显示集群的总带宽，而右轴显示每块硬盘的测量 I/O 带宽。WebDataset 和 AIStore 线性扩展到大约 300 个客户端，在此点之后，它们越来越多地受到来自旋转硬盘的最大 I/O 带宽的限制（每块硬盘约 150 MBytes/s）。为了比较，这里展示了 HDFS。HDFS 采用与 AIStore/WebDataset 类似的方法，并且也表现出线性扩展到大约 192 个客户端；在此点，它达到每块硬盘约 120 MBytes/s 的性能限制，并且当使用超过 1024 个客户端时失败。与 HDFS 不同，基于 WebDataset 的代码仅使用标准 URL 和 HTTP 来访问数据，并且与本地文件、存储在 Web 服务器上的文件以及 AIStore 工作方式相同。为了比较，在类似实验中，NFS 每块硬盘提供大约 10-20 MBytes/s。

存储数据集在 Tar 归档中

WebDataset 使用的格式是标准的 POSIX tar 归档，与用于备份和数据分发的归档相同。为了使用该格式存储用于深度学习的训练样本，我们采用了一些简单的命名约定：

  • 数据集是 POSIX tar 归档
  • 每个训练样本由具有相同基本名的相邻文件组成
  • 数据块按顺序编号

例如，ImageNet 存储在 1282 个独立的 100 兆字节数据块中，名称为 pythonimagenet-train-000000.tar to imagenet-train-001281.tar, 第一个数据块的内容是：

-r--r--r-- bigdata/bigdata      3 2020-05-08 21:23 n03991062_24866.cls
-r--r--r-- bigdata/bigdata 108611 2020-05-08 21:23 n03991062_24866.jpg
-r--r--r-- bigdata/bigdata      3 2020-05-08 21:23 n07749582_9506.cls
-r--r--r-- bigdata/bigdata 129044 2020-05-08 21:23 n07749582_9506.jpg
-r--r--r-- bigdata/bigdata      3 2020-05-08 21:23 n03425413_23604.cls
-r--r--r-- bigdata/bigdata 106255 2020-05-08 21:23 n03425413_23604.jpg
-r--r--r-- bigdata/bigdata      3 2020-05-08 21:23 n02795169_27274.cls

WebDataset 数据集可以直接从本地磁盘、网络服务器（因此得名）、云存储和对象存储中使用，只需更改 URL。WebDataset 数据集可以直接用于训练，无需解包，甚至可以在流式数据上执行训练，无需本地存储。

在许多深度学习应用中，训练过程中的洗牌非常重要，WebDataset 在分块级别和样本级别都执行洗牌。跨多个工作者的数据分割是在分块级别进行的，使用用户提供的 shard_selection 函数，默认为基于 get_worker_info. 进行分割的函数（WebDataset 可以与 tensorcom 库结合使用，以卸载解压缩/数据增强，并提供 RDMA 和直接到 GPU 的加载；详见下文。）

代码示例

下面是一些代码片段，展示了在典型的 PyTorch 深度学习应用中使用 WebDataset 的方法（完整示例可在 http://github.com/tmbdev/pytorch-imagenet-wds 找到。）

import webdataset as wds
import ...

sharedurl = "/imagenet/imagenet-train-{000000..001281}.tar"

normalize = transforms.Normalize(
  mean=[0.485, 0.456, 0.406],
  std=[0.229, 0.224, 0.225])

preproc = transforms.Compose([
  transforms.RandomResizedCrop(224),
  transforms.RandomHorizontalFlip(),
  transforms.ToTensor(),
  normalize,
])

dataset = (
  wds.Dataset(sharedurl)
  .shuffle(1000)
  .decode("pil")
  .rename(image="jpg;png", data="json")
  .map_dict(image=preproc)
  .to_tuple("image", "data")
)

loader = torch.utils.data.DataLoader(dataset, batch_size=64, num_workers=8)

for inputs, targets in loader:
  ...

这段代码几乎与 PyTorch Imagenet 示例中的基于文件的 I/O 管道相同：它创建了一个预处理/增强管道，使用该管道和数据源位置实例化了一个数据集，然后从数据集中构建了一个 DataLoader 实例。

WebDataset 使用流畅的 API 进行配置，内部构建处理管道。在不添加任何处理阶段的情况下，在这个例子中，WebDataset 与 PyTorch DataLoader 类一起使用，该类在多个线程中复制 DataSet 实例，并执行并行 I/O 和并行数据增强。

WebDataset 实例本身只是迭代每个训练样本作为一个字典：

# load from a web server using a separate client process
sharedurl = "pipe:curl -s http://server/imagenet/imagenet-train-{000000..001281}.tar"

dataset = wds.Dataset(sharedurl)

for sample in dataset:
  # sample["jpg"] contains the raw image data
  # sample["cls"] contains the class
  ...

想要了解我们如何使用 WebDataset 进行大规模训练的通用介绍，请查看这些 YouTube 视频。

  • AIStore 是一个开源的对象存储，能够实现全带宽的磁盘到 GPU 的数据传输（这意味着如果你有 1000 个转速为 200 MB/s 的旋转硬盘，AIStore 实际上可以向 GPU 提供 200 GB/s 的总带宽）。AIStore 完全兼容 WebDataset 客户端，并且理解 WebDataset 格式，允许它在存储系统中直接执行洗牌、排序、ETL 以及一些 map-reduce 操作。AIStore 可以被视为分布式对象存储、网络文件系统、分布式数据库和 GPU 加速的 map-reduce 实现的混合体。

  • tarp 是一个用于分割、合并、洗牌和处理 tar 存档和 WebDataset 数据集的小型命令行程序。

  • tensorcom 是一个支持分布式数据增强和 RDMA 到 GPU 的库。

  • pytorch-imagenet-wds 包含了一个如何使用 WebDataset 与 ImageNet 结合的示例，基于 PyTorch ImageNet 示例。

  • 大数据 2019 论文及基准

查看该库并提供您对 RFC 38419 的反馈。