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
Dependency:
- Keras 2.0
- Tensorflow
- OpenCV
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