如何使用PyTorch进行多进程处理？

Question

我正在尝试在pytorch中使用python的多进程Pool方法来处理图像。以下是代码：

from multiprocessing import Process, Pool
from torch.autograd import Variable
import numpy as np
from scipy.ndimage import zoom

def get_pred(args):

  img = args[0]
  scale = args[1]
  scales = args[2]
  img_scale = zoom(img.numpy(),
                     (1., 1., scale, scale),
                     order=1,
                     prefilter=False,
                     mode='nearest')

  # feed input data
  input_img = Variable(torch.from_numpy(img_scale),
                     volatile=True).cuda()
  return input_img

scales = [1,2,3,4,5]
scale_list = []
for scale in scales: 
    scale_list.append([img,scale,scales])
multi_pool = Pool(processes=5)
predictions = multi_pool.map(get_pred,scale_list)
multi_pool.close() 
multi_pool.join()

我遇到了这个错误：

`RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method

在这行中：

predictions = multi_pool.map(get_pred,scale_list)

有谁能告诉我我做错了什么吗？

- Rahul

Answer 1

根据 pytorch 文档，处理多进程的最佳实践是使用 torch.multiprocessing 而不是 multiprocessing。

要注意的是，在 Python 3 中，只有通过 spawn 或 forkserver 作为启动方法才支持在进程之间共享 CUDA 张量。

在不修改代码的情况下，您可以采用以下方法解决错误：替换

from multiprocessing import Process, Pool

随着：

from torch.multiprocessing import Pool, Process, set_start_method
try:
     set_start_method('spawn')
except RuntimeError:
    pass

- nicobonne

---

有时甚至需要使用torch.multiprocessing.set_start_method('spawn', force=True)。 - mr_mo

请确保将你的主循环用if __name__ == '__main__':包围起来，因为全局语句会在生成子进程时运行。 - xjcl

这解决了一个问题但又引入了另一个问题，TypeError: 'NoneType' object is not callable。有人见过这个吗？ - Mikhail

Answer 2

我建议您阅读多进程模块的文档，特别是此部分。您需要通过调用set_start_method来更改子进程创建方式。摘自上述引用文档：

import multiprocessing as mp

def foo(q):
    q.put('hello')

if __name__ == '__main__':
    mp.set_start_method('spawn')
    q = mp.Queue()
    p = mp.Process(target=foo, args=(q,))
    p.start()
    print(q.get())
    p.join()

- Oliver Baumann