在不同 CPU 上并行运行的多个 PyTorch 网络答案

【问题标题】：Multiple PyTorch networks running in parallel on different CPUs在不同 CPU 上并行运行的多个 PyTorch 网络
【发布时间】：2020-06-06 17:42:46
【问题描述】：

我试图让不同的 PyTorch 神经网络在不同的 CPU 上并行运行，但我发现与顺序运行它们相比，它并没有带来任何形式的加速。

下面是我的代码，它准确地复制了这个问题。如果您运行此代码，则表明使用 2 个进程所需的时间大约是使用 1 个进程运行的时间的两倍，但实际上应该花费相同的时间。

import time
import torch.multiprocessing as mp
import gym
import numpy as np
import copy
import torch.nn as nn
import torch

class NN(nn.Module):
    def __init__(self, output_dim):
        nn.Module.__init__(self)
        self.fc1 = nn.Linear(4, 50)
        self.fc2 = nn.Linear(50, 500)
        self.fc3 = nn.Linear(500, 5000)
        self.fc4 = nn.Linear(5000, output_dim)
        self.relu = nn.ReLU()

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.relu(self.fc2(x))
        x = self.relu(self.fc3(x))
        x = self.fc4(x)
        return x

def Worker(ix):
  print("Starting training for worker ", ix)
  env = gym.make('CartPole-v0')
  model = NN(2)
  for _ in range(2000):
    model(torch.Tensor(env.reset()))
  print("Finishing training for worker ", ix)

def overall_process(num_workers):
  workers = []
  for ix in range(num_workers):
    worker = mp.Process(target=Worker, args=(ix, ))
    workers.append(worker)
  [w.start() for w in workers]  
  for worker in workers:
    worker.join()

  print("Finished Training")  
  print(" ")

start = time.time()
overall_process(1)
print("Time taken: ", time.time() - start)
print(" ")

start = time.time()
overall_process(2)
print("Time taken: ", time.time() - start)

有谁知道为什么会发生这种情况以及如何解决？

我认为这可能是因为 PyTorch 网络会在后台自动实现 CPU 并行性，所以我尝试添加以下 2 行，但并不总能解决问题：

torch.set_num_threads(1)
torch.set_num_interop_threads(1)

【问题讨论】：

多少个核心？ AFAIK PyTorch 通过 MKL 使用所有可用的核心，因此一个网络使用所有核心的训练速度可能大约是两倍，这可以解释你的结果。尝试在不使用 torch.multiprocessing 的情况下使用 torch.set_num_threads(cores // 2)，看看这是否有任何改变。
不幸的是，做 torch.set_num_threads(1) 似乎并没有解决问题。

标签： python parallel-processing pytorch

【解决方案1】：

答案是在每个工作进程的开头（而不是在主进程中）设置 torch.set_num_threads(1)，如下所述：https://discuss.pytorch.org/t/multiple-networks-running-in-parallel-on-different-cpus/70482

【讨论】：