在 PyTorch 中真正并行运行“nn.ModuleList”中的独立“nn.Module”实例

Question

我有一个 PyTorch 模型，由存储在

FullyConnectedNetwork

内的多个独立

nn.ModuleList

实例组成。这是代码：

import torch.nn as nn

class FullyConnectedNetwork(nn.Module):
    def __init__(self):
        super(FullyConnectedNetwork, self).__init__()
        self.fc1 = nn.Linear(20, 10)
        self.fc2 = nn.Linear(10, 1)
    
    def forward(self, x):
        x = self.fc1(x)
        x = self.fc2(x)
        return x

class ParallelFCN(nn.Module):
    def __init__(self, n):
        super(ParallelFCN, self).__init__()
        self.models = nn.ModuleList([FullyConnectedNetwork() for _ in range(n)])
    
    def forward(self, x):
        outputs = [model(x[:, i*20:(i+1)*20]) for i, model in enumerate(self.models)]
        return torch.cat(outputs, dim=1)

# Example usage:
n = 1000
model = ParallelFCN(n)
print(model)

目前，我正在使用 for 循环通过每个

FullyConnectedNetwork

实例传递数据。然而，我意识到这种方法在软件意义上并不是真正的并行。

鉴于每个

FullyConnectedNetwork

都是独立的，有没有办法真正并行运行它们，也许使用多线程、多处理或 PyTorch 中的任何其他方法？

我需要它，因为我的模块数量可能非常大，大到 400 个，然后使用 for 循环进行处理非常慢。

Answer 1

我假设它在某种程度上是由 Cuda 处理的 - 它尝试在流中以异步方式执行内核，因此如果可能的话，内核调用可能会在时间上重叠。据我所知，Pytorch 中没有明确的 API 可以自行控制。

参考资料：

• Pytorch论坛相关讨论
• nn.Parallel 功能请求