2016-05-19 2 views
2

У меня есть нейронная сеть Caffe, и я хочу сделать переадресацию (используя GPU) в сети, не блокируя основной поток. Я использую python, я пытался использовать потоки и многопроцессорность. Они используют процессор, хотя я прошу его использовать графический процессор. Вот мой код:Перемещение вперед в Caffe NN параллельно

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 

ответ

2

У меня была очень похожая проблема.

использованием

caffe.set_mode_gpu() 

у ребенка нить сам решил.

Так попробуйте:

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)