如何通过在多个 GPU 上运行的多个衍生进程以及每个处理的数据批次以与输入相同的顺序获得输出？

Question

我正在使用 Pytorch 分布式数据并行方法并生成多个进程，每个进程在单独的 GPU 上运行。我正在使用 Pytorch 分布式数据采样器和数据加载器来将输入数据批量加载到每个进程。 我的问题：

1. 在引擎盖下，Pytorch 分布式数据采样器、数据加载器如何对输入数据进行分片？ 为简单起见，假设我们有 4 个 GPU，以及 400 个输入样本和批量大小，例如 50，然后 Pytorch 分布式数据采样器（连同数据加载器）将前 50 个样本 go 发送到 GPU-0，接下来的 50 个样本发送到 GPU-1。下一个 50 到 GPU-2，然后是 GPU-3，然后再下一个 50 到 GPU-0，即按照 GPU 设备编号的顺序？ 或者下一批输入的 GPU 到 select 的顺序是随机的，基于哪个 GPU 首先完成了上一批？ 还是先将 400 个样本分成 4 个部分，然后 GPU-0 将获得前 100 个输入数据样本（一次 50 个），GPU-1 将获得接下来的 100 个样本（一次 50 个）等等。 .在这种情况下，无论 GPU-3 是否比 GPU-0 更早开始其第二批，但仍然就输入数据而言，GPU-0 仍将拥有前 100 个样本，而 GPU-3 将拥有最后 100 个样本？

2). 我的第二个问题是如何以与输入数据相同的顺序检索 output 数据，以便最终合并的 output（将所有进程的输出组合在一个数据结构中）与原始输入的顺序相同，并且每个 output 对应于正确的输入

Answer 1

1. DistributedSampler上的 PyTorch 文档不提供关于数据如何跨进程和设备分布的任何保证，除了它实际上是跨进程和设备分布的事实。 您不应将应用程序设计为依赖于外部 package 的实现细节； 否则，您的应用程序可能会在更新 PyTorch 一天后突然失败，并且您不知道为什么（或者甚至可能无法开始）。 如果出于某种原因，您绝对需要以非常特定的方式分发数据，那么您应该推出自己的解决方案。 DistributedDataParallel的文档建议，如果您使用带有 N 个 GPU 的单个主机，则应该启动 N 个进程，每个进程指定一个 GPU。一个简单的解决方案是将进程的等级设置为等于指定的 GPU 设备 ID； 这反过来可以用于自定义采样器 class 到 select 适当的样本索引。
2. 您可以尝试控制各种分布式进程返回输出的顺序，但这会引入不必要的同步，这会破坏并行化的大部分目的。 更好的解决方案是以任意顺序简单地返回输出，然后在事后对它们进行排序。 如果您希望输出按与输入相同的顺序排序，您可以将每个输入与 integer 索引相关联（输入 0 获得索引 0，输入 1 获得索引 1，依此类推）。 返回 output 时，还返回关联输入的索引（例如作为元组）。 之后，您可以按相应的索引对输出进行排序。