Keras + Tensorflow: прогнозирование на нескольких графических процессорах

Я использую Keras с tensorflow в качестве бэкэнда. У меня есть одна скомпилированная/обученная модель.

мой цикл прогнозирования медленный, поэтому я хотел бы найти способ распараллелить predict_proba звонки, чтобы ускорить процесс. Я хотел бы взять список пакетов (данных), а затем на доступный GPU, запустите model.predict_proba() над подмножеством этих пакетов.
По существу:

data = [ batch_0, batch_1, ... , batch_N ]
on gpu_0 => return predict_proba(batch_0)
on gpu_1 => return predict_proba(batch_1)
...
on gpu_N => return predict_proba(batch_N)

Я знаю, что в чистом Tensorflow можно назначить ops для данного gpu (https://www.tensorflow.org/tutorials/using_gpu). Однако я не знаю, как это переводится в мою ситуацию, так как я построил/скомпилировал/обучил свою модель с помощью api Keras.

Я думал, что, возможно, мне просто нужно использовать модуль многопроцессорной обработки python и запустить процесс на gpu, который будет работать predict_proba(batch_n). Я знаю, что это теоретически возможно, учитывая еще один мой пост:Keras + Tensorflow и многопроцессорная обработка в Python. Однако, это все еще оставляет меня с дилеммой не зная, как на самом деле "выбрать" gpu для управления процессом.

мой вопрос сводится к: как можно распараллелить прогноз для одной модели в Keras через несколько графических процессоров при использовании Tensorflow в качестве бэкэнда Keras?

кроме того, мне интересно, возможно ли подобное распараллеливание для прогнозирования только с одним gpu.

высокоуровневое описание или пример кода были бы значительно ценю!

спасибо!

13
keras python python-multiprocessing tensorflow

2 ответов

Я создал один простой пример, чтобы показать, как запустить модель keras через несколько графических процессоров. В принципе, создается несколько процессов, и каждый из процессов владеет gpu. Чтобы указать идентификатор gpu в процессе, установка переменной env CUDA_VISIBLE_DEVICES является очень простым способом (os.environ ["CUDA_VISIBLE_DEVICES"]). Надеюсь, это git repo поможет вам.

https://github.com/yuanyuanli85/Keras-Multiple-Process-Prediction

4

вы можете использовать эту функцию для распараллеливания модели Keras (кредиты kuza55).
https://github.com/kuza55/keras-extras/blob/master/utils/multi_gpu.py
.

from keras.layers import merge
from keras.layers.core import Lambda
from keras.models import Model

import tensorflow as tf

def make_parallel(model, gpu_count):
    def get_slice(data, idx, parts):
        shape = tf.shape(data)
        size = tf.concat([ shape[:1] // parts, shape[1:] ],axis=0)
        stride = tf.concat([ shape[:1] // parts, shape[1:]*0 ],axis=0)
        start = stride * idx
        return tf.slice(data, start, size)

    outputs_all = []
    for i in range(len(model.outputs)):
        outputs_all.append([])

    #Place a copy of the model on each GPU, each getting a slice of the batch
    for i in range(gpu_count):
        with tf.device('/gpu:%d' % i):
            with tf.name_scope('tower_%d' % i) as scope:

                inputs = []
                #Slice each input into a piece for processing on this GPU
                for x in model.inputs:
                    input_shape = tuple(x.get_shape().as_list())[1:]
                    slice_n = Lambda(get_slice, output_shape=input_shape, arguments={'idx':i,'parts':gpu_count})(x)
                    inputs.append(slice_n)                

                outputs = model(inputs)

                if not isinstance(outputs, list):
                    outputs = [outputs]

                #Save all the outputs for merging back together later
                for l in range(len(outputs)):
                    outputs_all[l].append(outputs[l])

    # merge outputs on CPU
    with tf.device('/cpu:0'):
        merged = []
        for outputs in outputs_all:
            merged.append(merge(outputs, mode='concat', concat_axis=0))

        return Model(input=model.inputs, output=merged)
1