在多个分布式 CPU 上运行 PyTorch 数据加载器/数据集

Posted 2023-03-27

【中文标题】在多个分布式 CPU 上运行 PyTorch 数据加载器/数据集

【英文标题】：Running a PyTorch dataloader/Dataset on multiple distributed CPUs

【发布时间】：2020-04-20 10:57:10

【问题描述】：

我想知道是否有一种方法可以将数据加载器/数据集分发到多个 CPU，即使使用单个 GPU 也是如此。 具体来说，我想要一个 Dataset 类，__getitem__ 函数将分布在许多不同的 CPU 上（可能使用 mpi？但任何其他方式也不错）。

谢谢

编辑 我的标题被错误编辑了，我不是要分发模型本身，我只想分发模型的数据加载/解析

编辑 - 2 在这个方向上的一些有趣的讨论是可用的here

【问题讨论】：

DataLoader 在使用大于零的num_workers 初始化时的默认行为是在迭代时生成工作进程并并行加载数据

是的，我知道，但这受到计算机上可用 CPU 内核数量的限制。我有一个由 100 个 CPU 组成的集群，我想请所有这些 CPU 为我准备数据。

【参考方案1】：

您当然可以这样做，但请注意 - 由于通信成本高昂，它对于一般机器学习需求并不总是非常有效。 使用DistributedDataParallel

实现基于分布式数据并行 模块级别的 torch.distributed 包。

此容器通过以下方式并行化给定模块的应用程序 通过在指定设备中分块来拆分输入 批量维度。该模块在每台机器上复制 设备，并且每个这样的副本处理输入的一部分。期间 向后传递，每个节点的梯度被平均。

在实践中，我建议您使用pytorch_lightning 包，以减少您为此工作所需编写的一些样板代码。

参考： DistributedDataParallel,pytorch_lightning

【讨论】：

谢谢@Mano，我的标题被错误地编辑了。我不是要分发模型本身。我想分发数据加载/数据解析。谢谢