作者：Alex Aizman, Gavin Maltby, Thomas Breuel

数据集日益庞大，而 GPU 速度越来越快。这意味着深度学习研究人员和工程师有更多数据集可用于训练和验证其模型。

  • 许多用于静态图像识别研究的数据集正在涌现，包含 1000 万或更多图像，其中包括 OpenImages 和 Places。
  • 百万 YouTube 视频 (YouTube 8M) 在 720p 下消耗约 300 TB，用于物体识别、视频分析和动作识别研究。
  • Tobacco Corpus 数据集包含约 2000 万扫描高清页面，可用于 OCR 和文本分析研究。

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

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

  • 数据集大小：数据集通常超出节点本地磁盘存储容量，需要分布式存储系统和高效的网络访问。
  • 文件数量：数据集通常包含数十亿个文件，具有均匀随机访问模式，这常常会使本地文件系统和网络文件系统不堪重负。
  • 数据速率：在大型数据集上进行的训练任务通常使用多个 GPU，要求对数据集具有数 GB/秒的聚合 I/O 带宽；这只能通过大规模并行 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 带宽利用率。适用于 PB 级深度学习。
云计算 WebDataset 深度学习任务可以直接针对存储在云存储桶中的数据集进行训练；无需卷插件。本地任务和云任务的工作方式相同。适用于 PB 级学习。
配备现有分布式文件系统或对象存储的本地集群 WebDataset 的大顺序读取可以改善与现有分布式存储的性能，并消除了对专用卷插件的需求。
教育环境 WebDataset 可以存储在现有 Web 服务器和 Web 缓存中，学生可以通过 URL 直接访问
从本地驱动器在工作站上进行训练 数据仍在下载时，任务即可开始训练。训练数据无需解包。与基于文件的随机访问数据集相比，在硬盘驱动器上 I/O 性能提高了十倍。
所有环境 数据集以归档格式表示，包含文件类型等元数据。数据采用原生格式（JPEG、MP4 等）压缩。数据管理、ETL 式任务以及数据转换和 I/O 得到简化且易于并行化。

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

高性能

为了在本地集群上实现高性能计算，配套的开源 AIStore 服务器提供完整的磁盘到 GPU 数据传输带宽，仅受限于硬件约束。这篇 Bigdata 2019 论文包含详细的基准测试和性能测量结果。除了基准测试之外，NVIDIA 和 Microsoft 的研究项目也使用 WebDataset 处理 PB 级数据集和数十亿个训练样本。

以下是 AIStore 与 WebDataset 客户端的基准测试结果，使用了 12 个服务器节点，每个节点配备 10 个旋转硬盘驱动器。

左轴显示来自集群的总带宽，而右轴显示测得的每驱动器 I/O 带宽。WebDataset 和 AIStore 可以线性扩展到约 300 个客户端，此时它们越来越受限于旋转硬盘驱动器提供的最大 I/O 带宽（每个驱动器约 150 MB/秒）。作为比较，图中显示了 HDFS。HDFS 使用与 AIStore/WebDataset 类似的方法，并且也可以线性扩展到约 192 个客户端；此时，它达到每驱动器约 120 MB/秒的性能极限，并且在使用超过 1024 个客户端时发生故障。与 HDFS 不同，基于 WebDataset 的代码仅使用标准 URL 和 HTTP 访问数据，并且与本地文件、存储在 Web 服务器上的文件以及 AIStore 的工作方式相同。作为比较，在类似实验中，NFS 的每驱动器传输速率约为 10-20 MB/秒。

将数据集存储在 Tar 归档文件中

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

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

例如，ImageNet 存储在 1282 个独立的 100 MB 分片中，名称为 pythonimagenet-train-000000.tar 到 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 数据集可以直接从本地磁盘、Web 服务器（因此得名）、云存储和对象存储中使用，只需更改 URL 即可。WebDataset 数据集无需解包即可用于训练，甚至可以在流式数据上进行训练，无需本地存储。

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

代码示例

以下是一些代码片段，展示了 WebDataset 在典型 PyTorch 深度学习应用中的使用（完整示例可在 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 类一起使用，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/秒的旋转硬盘驱动器，AIStore 实际上可以向 GPU 提供总计 200 GB/秒 的带宽）。AIStore 作为客户端完全兼容 WebDataset，此外还理解 WebDataset 格式，使其能够直接在存储系统中执行混洗、排序、ETL 和一些 MapReduce 操作。AIStore 可以被视为分布式对象存储、网络文件系统、分布式数据库和 GPU 加速 MapReduce 实现的结合体。

  • tarp 是一个小型命令行程序，用于分割、合并、混洗和处理 tar 归档文件和 WebDataset 数据集。

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

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

  • 包含基准测试的 Bigdata 2019 论文

查阅该库，并就RFC 38419 提供你的反馈意见。