Simple Example to run Keras models in multiple processes
This git repo contains an example to illustrate how to run Keras models prediction in multiple processes with multiple gpus. Each process owns one gpu. I wanted to run prediction by using multiple gpus, but did not find a clear solution after searching online. So, I created this example to show how to do that. Hope this git repo can help others who met the same problem.
This software works as a producer-consumer model. The scheduler scans the image path and put all of them into a Queue; while each worker as a separate process process the images in the Queue. If all of images are proceeded, the worker process will terminate.
Pay attention to the implementation Vgg16Worker::Run()
def run(self): #set environment os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = str(self._gpuid) #load models import vgg16 #download the vgg16 weights from xnet = vgg16.Vgg16('vgg16_weights_tf_dim_ordering_tf_kernels.h5') print 'vggnet init done', self._gpuid while True: xfile = self._queue.get() if xfile == None: self._queue.put(None) break label = self.predict(xnet, xfile) print 'woker', self._gpuid, ' xfile ', xfile, " predicted as label", label
How to run the sample
- Keras 2.0
Download VGG16 Keras Model from: https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg16_weights_tf_dim_ordering_tf_kernels.h5
Command: usage: main.py [-h] [--imgpath IMGPATH] [--gpuids GPUIDS]
optional arguments: -h, --help show this help message and exit --imgpath IMGPATH path to your images to be proceed --gpuids GPUIDS gpu ids to run