python - Keras+Tensorflow: Predicción en múltiples gpus

python-multiprocessing

Question

Estoy usando Keras con tensorflow como back-end. Tengo un modelo compilado / entrenado.

Mi ciclo de predicción es lento, así que me gustaría encontrar una forma de paralelizar las llamadas a predict_proba para acelerar las cosas. Me gustaría tomar una lista de lotes (de datos) y luego, por gpu disponible, ejecutar model.predict_proba() sobre un subconjunto de esos lotes.
Esencialmente:

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) 

Sé que es posible en Tensorflow puro asignar operaciones a un gpu determinado ( https://www.tensorflow.org/tutorials/using_gpu ). Sin embargo, no sé cómo se traduce esto en mi situación, ya que he construido / compilado / entrenado mi modelo usando la API de Keras.

Pensé que tal vez solo necesitaba usar el módulo de multiprocesamiento de python y comenzar un proceso por gpu que ejecutaría predict_proba(batch_n) . Sé que esto es teóricamente posible dado otro post SO mío: Keras + Tensorflow y Multiprocesamiento en Python . Sin embargo, esto todavía me deja con el dilema de no saber cómo realmente "elegir" un gpu para operar el proceso.

Mi pregunta se reduce a: ¿cómo se puede paralelizar la predicción para un modelo en Keras a través de múltiples gpus cuando se usa Tensorflow como backend de Keras?

Además, tengo curiosidad si es posible una paralelización similar para la predicción con solo un gpu.

¡Una descripción de alto nivel o un ejemplo de código sería muy apreciado!

¡Gracias!

Answer 1

Creé un ejemplo simple para mostrar cómo ejecutar el modelo de kera en múltiples gpus. Básicamente, se crean múltiples procesos y cada uno de los procesos posee un gpu. Para especificar el id. De gpu en proceso, la configuración de la variable de entorno CUDA_VISIBLE_DEVICES es muy sencilla (os.environ ["CUDA_VISIBLE_DEVICES"]). Espero que este git repo pueda ayudarte.

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

Answer 2

Puede usar esta función para paralelizar un modelo de Keras (créditos a 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)