跳转到主要内容
博客

用于大型数据集、多文件、多 GPU 的高效 PyTorch I/O 库

作者： 2020 年 8 月 11 日2024 年 11 月 16 日暂无评论

数据集日益庞大，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 数据传输（这意味着如果您有 1000 个旋转驱动器，每个读取速度为 200 MB/s，AIStore 实际上可以为 GPU 提供 200 GB/s 的聚合带宽）。这篇 Bigdata 2019 论文包含详细的基准和性能测量。除了基准测试之外，NVIDIA 和微软的研究项目已使用 WebDataset 处理 PB 级数据集和数十亿个训练样本。

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

左轴显示集群的聚合带宽，右轴显示测得的每驱动器 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 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 类一起使用，该类在多个线程中复制 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 和一些 MapReduce 操作。AIStore 可以被视为分布式对象存储、网络文件系统、分布式数据库和 GPU 加速 MapReduce 实现的混合体。
  • tarp 是一个小型命令行程序，用于分割、合并、打乱和处理 tar 归档文件和 WebDataset 数据集。
  • tensorcom 是一个支持分布式数据增强和 RDMA 到 GPU 的库。
  • pytorch-imagenet-wds 包含一个如何使用 WebDataset 和 ImageNet 的示例，基于 PyTorch ImageNet 示例。
  • 附有基准测试的 Bigdata 2019 论文

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