Multiprocessing using cuda

Viewed 6331

I am trying to run multiprocessing in my python program. I created two processes and passed a neural network in the one process and some heavy computational function in the other. I wanted the neural net to run on GPU and the other function on CPU and thereby I defined neural net using cuda() method.

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

But when I run the program the got the following error: RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the ‘spawn’ start method

So I tried with spawn as well as forkserver start method, but then I got the other error: RuntimeError: cuda runtime error (71) : operation not supported at …/torch/csrc/generic/StorageSharing.cpp:245

I have tried python3 multiprocessing and torch.multiprocessing both but nothing worked for me.

1 Answers

put set_start_method inside if __name__ == '__main__':

the code will be

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()
Related