时间戳

首页 > 解决方案 > PyTorch：如何使用 multiprocessing.pool 在多个 GPU 上并行化

python - PyTorch：如何使用 multiprocessing.pool 在多个 GPU 上并行化

问题描述

我有以下代码，我试图在 PyTorch 中的多个 GPU 上并行化：

import numpy as np
import torch
from torch.multiprocessing import Pool

X = np.array([[1, 3, 2, 3], [2, 3, 5, 6], [1, 2, 3, 4]])
X = torch.DoubleTensor(X).cuda()

def X_power_func(j):
    X_power = X**j
    return X_power

if __name__ == '__main__':
  with Pool(processes = 2) as p:   # Parallelizing over 2 GPUs
    results = p.map(X_power_func, range(4))

results

但是当我运行代码时，我收到了这个错误：

---------------------------------------------------------------------------
RemoteTraceback                           Traceback (most recent call last)
RemoteTraceback: 
"""
Traceback (most recent call last):
  File "/usr/lib/python3.6/multiprocessing/pool.py", line 119, in worker
    result = (True, func(*args, **kwds))
  File "/usr/lib/python3.6/multiprocessing/pool.py", line 44, in mapstar
    return list(map(*args))
  File "<ipython-input-35-6529ab6dac60>", line 11, in X_power_func
    X_power = X**j
RuntimeError: CUDA error: initialization error
"""

The above exception was the direct cause of the following exception:

RuntimeError                              Traceback (most recent call last)
<ipython-input-35-6529ab6dac60> in <module>()
     14 if __name__ == '__main__':
     15   with Pool(processes = 1) as p:
---> 16     results = p.map(X_power_func, range(8))
     17 
     18 results

1 frames
/usr/lib/python3.6/multiprocessing/pool.py in get(self, timeout)
    642             return self._value
    643         else:
--> 644             raise self._value
    645 
    646     def _set(self, i, obj):

RuntimeError: CUDA error: initialization error

我哪里出错了？任何帮助将不胜感激。

标签： pythonmultiprocessingpytorchgpu

解决方案

我认为通常的方法是在多处理之前调用model.share_memory()一次，假设你有一个子类模型nn.Module。对于张量，它应该是X.share_memory_(). 不幸的是，我很难让它与您的代码一起使用，如果X.share_memory_()在调用 pool.map 之前调用它会挂起（没有错误）；我不确定原因是否是因为 X 是一个全局变量，它没有作为 map 中的参数之一传递。

什么工作是这样的：

X = torch.DoubleTensor(X)

def X_power_func(j):
    X_power = X.cuda()**j
    return X_power

顺便说一句：https ://github.com/pytorch/pytorch/issues/15734提到“在你分叉之前不能初始化 CUDA API ”（这可能是你看到的问题）。

另外https://github.com/pytorch/pytorch/issues/17680如果在 Jupyter 笔记本中使用 spawn “spawn 方法将运行笔记本顶层中的所有内容”（可能是我在代码挂起时看到的问题，在笔记本）。简而言之，我无法让 fork 或 spawn 工作，除非使用上面的序列（在分叉进程中之前不使用 CUDA）。

推荐阅读