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 elpredict_proba llama a acelerar las cosas. Me gustaría tomar una lista de lotes (de datos) y luego por gpu disponible, ejecutarmodel.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 una gpu dada (https://www.tensorflow.org/tutorials/using_gpu) Sin embargo, no sé cómo esto se traduce 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 e iniciar un proceso por gpu que se ejecutepredict_proba(batch_n). Sé que esto es teóricamente posible dado otro puesto mío de SO:Keras + Tensorflow y Multiprocesamiento en Python. Sin embargo, esto todavía me deja con el dilema de no saber cómo "elegir" realmente una gpu para operar el proceso.

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

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

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

¡Gracias!