使用 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()

但是当我运行该程序时，出现以下错误：运行时错误：无法在分叉子进程中重新初始化 CUDA。要将 CUDA 与多处理一起使用，您必须使用"生成"启动方法

所以我尝试使用spawn和forkserver启动方法，但后来我得到了另一个错误： 运行时错误： cuda 运行时错误 (71( ： .../torch/csrc/generic/StorageSharing.cpp：245 不支持操作

我已经尝试了python3多处理和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()