python-3.x - 使用cuda进行多处理

Question

我正在尝试在我的 python 程序中运行多处理。我创建了两个进程，并在一个进程中传递了一个神经网络，在另一个进程中传递了一些繁重的计算函数。我希望神经网络在 GPU 上运行，而其他功能在 CPU 上运行，因此我使用 cuda() 方法定义了神经网络。

import cv2
import torch
import torch.nn as nn
import multiprocessing as mp    #I even tried import torch.multiprocessing
from multiprocessing import set_start_method
try:
    set_start_method('spawn')
except RunTimeError:
    pass

class network(nn.Module):
    '''neural net class'''

def func1(img, net):
    '''used to call neural net'''

def func2(img):
    '''image processing part'''

if __name__ == '__main__':
    net = network().cuda() #here network is my neural net class
    img = cv2.imread('img.png')
    p1 mp.Process(target= func1, args=(img, net))    #func1 is used to run neural net
    p2 = mp.Process(target=func2, args=(img,))    #func2 is used for some img-processing
    p1.start()
    p2.start()
    p1.join()
    p2.join()

但是当我运行程序时出现了以下错误:RuntimeError:无法在 fork 的子进程中重新初始化 CUDA。要将 CUDA 与多处理一起使用，您必须使用“spawn”启动方法

所以我尝试了 spawn 以及 forkserver 启动方法，但后来我得到了另一个错误:RuntimeError: cuda runtime error (71) : 操作不支持…/torch/csrc/generic/StorageSharing.cpp:245

我已经尝试过 python3 multiprocessing 和 torch.multiprocessing 两者，但对我没有任何效果。

Answer 1

将 set_start_method 放入 if __name__ == '__main__':

代码将是

import cv2
import torch
import torch.nn as nn
import multiprocessing as mp    #I even tried import torch.multiprocessing
from multiprocessing import set_start_method


class network(nn.Module):
    '''neural net class'''

def func1(img, net):
    '''used to call neural net'''

def func2(img):
    '''image processing part'''

if __name__ == '__main__':
    set_start_method('spawn')
    net = network().cuda() #here network is my neural net class
    img = cv2.imread('img.png')
    p1 mp.Process(target= func1, args=(img, net))    #func1 is used to run neural net
    p2 = mp.Process(target=func2, args=(img,))    #func2 is used for some img-processing
    p1.start()
    p2.start()
    p1.join()
    p2.join()