跳转到主要内容
博客

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

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

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

  • 许多用于静态图像识别研究的数据集正在不断涌现,包含1000万或更多图像,其中包括OpenImages和Places。
  • 数百万YouTube视频 (YouTube 8M) 在720p分辨率下消耗约300 TB,用于目标识别、视频分析和动作识别研究。
  • Tobacco Corpus包含约2000万页扫描的HD页面,对OCR和文本分析研究很有用。

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

然而,处理大量数据集带来了许多挑战

  • 数据集大小: 数据集通常超出节点本地磁盘存储的容量,需要分布式存储系统和高效的网络访问。
  • 文件数量: 数据集通常由数十亿个文件组成,具有均匀的随机访问模式,这常常会使本地和网络文件系统不堪重负。
  • 数据速率: 在大型数据集上进行的训练作业通常使用多个GPU,需要对数据集的总I/O带宽达到每秒数GB;这只能通过大规模并行I/O系统来满足。
  • 洗牌和增强: 训练数据需要在训练前进行洗牌和增强。
  • 可伸缩性: 用户通常希望在小数据集上开发和测试,然后迅速扩展到大型数据集。

传统的本地和网络文件系统,甚至对象存储服务器,都不是为这类应用设计的。 WebDataset I/O库 用于PyTorch,以及可选的 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直接访问。
工作站本地驱动器训练作业可以在数据仍在下载时开始训练。数据无需解压即可进行训练。与基于随机访问文件的D数据集相比,硬盘驱动器上的I/O性能提高了十倍。
所有环境数据集以归档格式表示,并包含文件类型等元数据。数据以本机格式(JPEG、MP4等)进行压缩。数据管理、ETL样式作业以及数据转换和I/O得到简化并易于并行化。

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

高性能

对于本地集群上的高性能计算,配套的开源 AIStore 服务器提供全带宽的磁盘到GPU数据传输(这意味着如果您有1000个转速硬盘,每个读取速度为200 MB/s,AIStore实际上可以向GPU提供200 GB/s的总带宽)。 这篇2019年大数据论文 包含详细的基准测试和性能测量。除了基准测试,NVIDIA和Microsoft的研究项目已将WebDataset用于PB级数据集和数十亿个训练样本。

以下是AIStore与WebDataset客户端的基准测试,使用12个服务器节点,每个节点有10个机械硬盘。

左轴显示集群的总带宽,右轴显示测得的每驱动器I/O带宽。WebDataset和AIStore线性扩展到大约300个客户端,此时它们越来越受到机械硬盘最大I/O带宽的限制(每驱动器大约150 MB/s)。为了进行比较,图中显示了HDFS。HDFS使用与AIStore/WebDataset类似的方法,也表现出线性扩展,直到大约192个客户端;此时,它达到了每驱动器约120 MB/s的性能限制,并且在使用超过1024个客户端时失败。与HDFS不同,基于WebDataset的代码仅使用标准URL和HTTP访问数据,并且与本地文件、存储在Web服务器上的文件以及AIStore的工作方式相同。相比之下,在类似的实验中,NFS每驱动器提供大约10-20 MB/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在分片级别和样本级别都执行洗牌。数据在多个worker之间的分割在分片级别进行,使用用户提供的 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进行配置,该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 提供您的反馈。