训练大型深度学习模型需要大量数据集。亚马逊简单存储服务（Amazon S3）是一种可扩展的云对象存储服务，用于存储大量训练数据集。机器学习（ML）从业者需要一个高效的数据管道，可以从亚马逊 S3 下载数据，转换数据，并将数据馈送到 GPU 进行高吞吐量、低延迟的训练模型。
在本文中，我们介绍了新的 S3 IO DataPipes for PyTorch， S3FileLister 和 S3FileLoader 。为了提高内存效率和运行速度，新的 DataPipes 使用 C++扩展来访问亚马逊 S3。基准测试显示， S3FileLoader 在从亚马逊 S3 下载自然语言处理（NLP）数据集方面比 FSSpecFileOpener 快 59.8%。您可以使用新的 DataPipes 构建 IterDataPipe 训练管道。我们还演示了新的 DataPipe 可以将 Bert 和 ResNet50 的整体训练时间缩短 7%。新的 DataPipes 已集成到开源 TorchData 0.4.0 中，与 PyTorch 1.12.0 版本一起。
概述
亚马逊 S3 是一个可扩展的云存储服务，没有数据量限制。从亚马逊 S3 加载数据并馈送到高性能 GPU（如 NVIDIA A100）可能具有挑战性。这需要能够满足 GPU 数据处理速度的高效数据管道。为了帮助解决这个问题，我们发布了一个新的高性能 PyTorch 工具：S3 IO DataPipes。DataPipes 是从 torchdata.datapipes.iter.IterDataPipe 派生的子类，因此它们可以与 IterableDataPipe 接口交互。开发者可以快速构建他们的 DataPipe DAGs，以访问、转换和操作数据，具有洗牌、分片和批量功能。
新的 DataPipes 被设计为对文件格式不可知，亚马逊 S3 数据以二进制大对象（BLOBs）的形式下载。它可以作为一个可组合的构建块来组装一个 DataPipe 图，该图可以将表格、NLP 和计算机视觉（CV）数据加载到您的训练管道中。
在内部，新的 S3 IO DataPipes 采用了 AWS C++ SDK 的 C++ S3 处理器。一般来说，与 Python 相比，C++实现更节省内存，并且在多线程中有更好的 CPU 核心利用率（没有全局解释器锁）。新的 C++ S3 IO DataPipes 推荐用于在训练大型深度学习模型时进行高吞吐量、低延迟的数据加载。
新的 S3 IO DataPipes 提供了两个一等公民 API：
- S3FileLister – 可迭代的对象，用于列出给定 S3 前缀内的 S3 文件 URL。此 API 的功能名称为
list_files_by_s3。 - S3FileLoader – 可迭代的对象，用于从给定的 S3 前缀加载 S3 文件。此 API 的功能名称为
load_files_by_s3。
使用说明
在本节中，我们提供使用新 S3 IO DataPipes 的说明。同时，我们还提供了一个代码片段用于 load_files_by_s3() 。
从源码构建
新的 S3 IO DataPipes 使用 C++扩展。它默认包含在 torchdata 包中。然而，如果新 DataPipes 在环境中不可用，例如 Windows 上的 Conda，您需要从源码构建。更多信息请参考 Iterable Datapipes。
配置
亚马逊 S3 支持全局存储桶。然而，存储桶是在一个区域内创建的。您可以通过使用 __init__() 将区域传递给 DataPipes。或者，您可以在 shell 中 export AWS_REGION=us-west-2 ，或者在代码中使用 os.environ['AWS_REGION'] = 'us-east-1' 设置环境变量。
要读取非公开访问的存储桶中的对象，您必须通过以下方法之一提供 AWS 凭证：
- 使用
AWS configure安装和配置 AWS 命令行界面（AWS CLI）。 - 在本地系统上的 AWS 凭证配置文件中设置凭证，该文件位于
~/.aws/credentials（Linux、macOS 或 Unix）。 - 设置
AWS_ACCESS_KEY_ID和AWS_SECRET_ACCESS_KEY环境变量 - 如果您在 Amazon Elastic Compute Cloud (Amazon EC2)实例上使用此库，请指定 AWS 身份和访问管理(IAM)角色，然后授予 EC2 实例对该角色的访问权限
示例代码
以下代码片段提供了 load_files_by_s3() 的典型用法：
from torch.utils.data import DataLoader
from torchdata.datapipes.iter import IterableWrapper
s3_shard_urls = IterableWrapper(["s3://bucket/prefix/",])
.list_files_by_s3()
s3_shards = s3_shard_urls.load_files_by_s3()
# text data
training_data = s3_shards.readlines(return_path=False)
data_loader = DataLoader(
training_data,
batch_size=batch_size,
num_workers=num_workers,
)
# training loop
for epoch in range(epochs):
# training step
for bach_data in data_loader:
# forward pass, backward pass, model update
基准
在本节中，我们展示了新的 DataPipe 如何减少 Bert 和 ResNet50 的整体训练时间。
针对 FSSpec 的独立 DataLoader 性能评估
FSSpecFileOpener 是另一个 PyTorch S3 DataPipe。它使用 botocore 和 aiohttp/asyncio 来访问 S3 数据。以下为性能测试设置和结果（摘自基于原生 AWSSDK 和 FSSpec（boto3）基于的 DataPipe 的性能比较）。
测试中的 S3 数据是一个分片文本数据集。每个分片大约有 10 万行，每行大约 1.6 KB，因此每个分片大约为 156 MB。该基准测试中的测量值是平均 1000 个批次的结果。没有进行洗牌、采样或转换。
下面的图表报告了各种批量大小的吞吐量比较，对于 num_workers=0 ，数据加载器在主进程中运行。 S3FileLoader 的每秒查询数（QPS）比 fsspec 高 90%，在批大小为 512 时。
下面的图表报告了 num_workers=4 的结果，数据加载器在主进程中运行。 S3FileLoader 在批大小为 512 时比 fsspec 高 59.8%。
针对 Boto3 训练 ResNet50 模型
对于以下图表，我们在 4 个 p3.16xlarge 实例的集群上训练了一个 ResNet50 模型，总共有 32 个 GPU。训练数据集是 ImageNet，包含 1.2 百万张图片，组织成 1,000 个 1,000 张图片的块。训练批次大小为 64。训练时间以秒为单位。对于 8 个 epoch， S3FileLoader 比 Boto3 快 7.5%。
针对 Boto3 训练 Bert 模型
对于以下批次，我们在 4 个 p3.16xlarge 实例的集群上训练了一个 Bert 模型，总共有 32 个 GPU。训练语料库有 1474 个文件。每个文件大约有 150,000 个样本。为了运行更短的 epoch，我们每个文件使用 0.05%（大约 75 个样本）。批次大小为 2,048。训练时间以秒为单位。对于 1 个 epoch， S3FileLoader 比 Boto3 快 7%。
与原始 PyTorch S3 插件进行比较
新的 PyTorch S3 DataPipes 在性能上显著优于原始的 PyTorch S3 插件。我们已经调整了内部缓冲区大小为 S3FileLoader 。加载时间以秒为单位进行测量。
在我们的实验中，对于 10 个分片化的 charades 文件（每个文件大约 1.5 GiB）， S3FileLoader 的速度快了 3.5 倍。
最佳实践
训练大型深度学习模型可能需要一个拥有数十甚至数百个节点的巨大计算集群。集群中的每个节点可能会生成大量针对特定 S3 分片的数据加载请求。为了避免限制，我们建议在 S3 桶和 S3 文件夹中分片训练数据。
为了实现良好的性能，文件大小足够大以并行处理给定文件是有帮助的，但又不至于过大，以免我们触及到在 Amazon S3 上该对象的吞吐量限制，这取决于训练作业。最佳大小可能在 50-200 MB 之间。
结论和下一步
在本文中，我们向您介绍了新的 PyTorch IO DataPipes。新的 DataPipes 使用 aws-sdk-cpp ，与基于 Boto3 的数据加载器相比，性能更佳。
对于下一步，我们计划通过关注以下功能来提高可用性、性能和功能：
- 使用 IAM 角色的 S3 授权 - 目前，S3 DataPipes 支持显式访问凭证、实例配置文件和 S3 存储桶策略。然而，在某些用例中，IAM 角色更受欢迎。
- 双缓冲 - 我们计划提供双缓冲以支持多工作线程下载。
- 本地缓存 - 我们计划使模型训练能够遍历训练数据集进行多次遍历。在第一个 epoch 之后进行本地缓存可以减少从 Amazon S3 的飞行延迟时间，这可以显著加速后续 epochs 的数据检索时间。
- 可定制配置 - 我们计划公开更多参数，例如内部缓冲区大小、多部分块大小和执行器数量，并允许用户进一步调整数据加载效率。
- Amazon S3 上传 - 我们计划扩展 S3 DataPipes 以支持检查点上传。
- 与 fsspec 合并 –
fsspec在其他系统中也被使用，如torch.save()。我们可以将新的 S3 DataPipes 与fsspec集成，使其有更多用例。
致谢
我们要感谢来自亚马逊的 Vijay Rajakumar 和 Kiuk Chung，他们为 S3 Common RunTime 和 PyTorch DataLoader 提供了指导。我们还要感谢来自 Meta AI/ML 的 Erjia Guan、Kevin Tse、Vitaly Fedyunin、Mark Saroufim、Hamid Shojanazeri、Matthias Reso 和 Geeta Chauhan，以及来自 AWS 的 Joe Evans，他们审阅了博客和 GitHub PR。
