如何在一个模块内并行运行多个模块？

Question

我一直在阅读一篇论文，其中包含以下机器学习模型，我想在 PyTorch 中复制该模型

本质上，输入被分割成 n 个大小相等的向量，并且每个向量都被传递到一个单独的本地模型。然后将本地模型的所有输出连接起来并运行到下一层。 （x与我的问题无关，所以我们忽略它）

到目前为止我已经想出了这个：

class GlobalModel(torch.nn.Module):
    def __init__(self, n_local_models):
        super(GlobalModel, self).__init__()

        self.local_models = [LocalModel() for _ in range(n_local_models)]
        self.linear = torch.nn.Linear(100, 100)  
        self.activation = torch.nn.ReLU()

LocalModel 是其他一些 torch.nn.Module。线性层大小只是一个虚拟的，稍后我将使其随本地模型动态变化。

我的问题是，在连接它们并将它们传递给线性层和激活函数之前，如何最好地编写一个并行运行所有局部模型的forward()函数。因为我能想到的实现这一点的唯一方法是迭代本地模型列表并按顺序执行每个模型。但这看起来相当慢，我觉得应该有一个更优雅的解决方案。

Answer 1

要并行运行多个模块，您通常需要多个进程（和多个 GPU）。 Python 不太擅长多线程（由于 GIL）。

即使没有 Python，GPU 编程也经过高度优化，可并行化在张量上执行的大型单独操作，而不是并行执行多个单独的操作。要在单个 GPU 上并行执行单独的操作，您需要将它们发送到不同的 GPU 执行器（“流”）。

您没有链接到该论文，但有时也会将一个操作拆分为多个并行模块，只是为了更好地利用多个 GPU。

Answer 2

一些澄清问题。

您可以使用多少个 GPU？

本地模型是预先训练的还是在线训练的？如果它们是预先训练的，是否有必要通过它们进行反向传播？

GPU 确实不喜欢在单卡内进行并行处理。如果您可以访问多张卡，则可以考虑跨多张卡并行化本地模型。如果您没有足够的 GPU 来执行此操作，那么最好串联运行本地模型。具有卡共享的多模型并行处理将是一个巨大的痛苦，并且很可能比串行运行慢。

如果本地模型是预先训练的并且您不需要通过它们进行反向传播，您可以考虑构建一个特征提取管道，在训练之前计算本地模型输出。