Reputation: 4873
I have a Caffe neural network and I want to do forward pass (using the GPU) in the network without blocking the main thread. I am using python, I tried using threading and multiprocessing. They are using the CPU even though I am asking it to use the GPU. Here's my code:
class ThreadingWorker(threading.Thread):
def __init__(self):
super(ThreadingWorker,self).__init__()
param = config()
self.model = param.model
self.net = caffe.Net(self.model.deployFile, self.model.caffemodel, caffe.TEST)
def run(self):
input_data = np.random.rand(1,4,self.model.width,self.model.height)
start = time()
self.net.forward(data=input_data)
print 'Success, took %f seconds' % (time()-start)
class MultProcessingWorker(mp.Process):
def run(self):
param = config()
self.model = param.model
self.net = caffe.Net(self.model.deployFile, self.model.caffemodel, caffe.TEST)
input_data = np.random.rand(1,4,self.model.width,self.model.height)
start = time()
self.net.forward(data=input_data)
print 'Success, took %f seconds' % (time()-start)
class NormalWorker(object):
'''Using the main thread, no parallelization is being used here'''
def __init__(self):
param = config()
self.model = param.model
self.net = caffe.Net(self.model.deployFile, self.model.caffemodel, caffe.TEST)
def run(self):
input_data = np.random.rand(1,4,self.model.width,self.model.height)
start = time()
self.net.forward(data=input_data)
print 'Success, took %f seconds' % (time()-start)
p = NormalWorker()
p.run()
>> Success, took 0.34 seconds
thread = ThreadingWorker()
thread.start()
>> Success, took 3.54 seconds
p = MultProcessingWorker()
p.start()
>> Success, took 3.45 seconds
Upvotes: 2
Views: 831
Reputation: 121
I had a very similar Problem.
using
caffe.set_mode_gpu()
in the child thread itself solved it.
So try:
class MultProcessingWorker(mp.Process):
def run(self):
caffe.set_mode_gpu() # <--- should fix the problem
param = config()
self.model = param.model
self.net = caffe.Net(self.model.deployFile, self.model.caffemodel, caffe.TEST)
input_data = np.random.rand(1,4,self.model.width,self.model.height)
start = time()
self.net.forward(data=input_data)
print 'Success, took %f seconds' % (time()-start)
Upvotes: 2