web-dev-qa-db-ja.com

PyTorchマルチプロセッシングの使用方法は?

画像を処理するためにPoolのpythonのマルチプロセッシングpytorchメソッドを使用しようとしています。これがコードです:

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)

誰かが私が間違っていることを教えてもらえますか?

4
Rahul

pytorchドキュメント で述べたように、マルチプロセッシングを処理するためのベストプラクティスは、multiprocessingではなくtorch.multiprocessingを使用することです。

プロセス間でのCUDAテンソルの共有は、Python 3、spawnまたはforkserverを開始メソッドとして使用する場合にのみサポートされています。

コードに触れずに、発生したエラーの回避策は

from multiprocessing import Process, Pool

と:

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

マルチプロセッシングモジュールのドキュメント、特に このセクション を読むことをお勧めします。 set_start_methodを呼び出して、サブプロセスの作成方法を変更する必要があります。これらの引用されたドキュメントから引用:

import multiprocessing as mp

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

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