[解決済み] PyTorchのマルチプロセッシングを利用するには？

質問内容

Pythonのマルチプロセッシングを使いたいのですが、どうすればいいですか？ Pool メソッドを 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)

何が間違っているのか、どなたか教えていただけませんか？

解決方法を教えてください。

にあるように pytorch ドキュメント マルチプロセッシングを処理するためのベストプラクティスは torch.multiprocessing の代わりに 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