首页 文章

Keras Tensorflow:对多个gpus的预测

提问于
浏览
13

我正在使用带有张量流的Keras作为后端 . 我有一个编译/训练模型 .

我的预测循环很慢,所以我想找到一种方法来并行化 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) . 但是,我没有使用Keras的api Build /编译/训练我的模型 .

我曾经想过,也许我只需要使用python的多处理模块,并为每个运行 predict_proba(batch_n) 的gpu启动一个进程 . 我知道这在理论上是可能的,因为我的另一个SO帖子:Keras + Tensorflow and Multiprocessing in Python . 然而,这仍然让我陷入两难境地:不知道如何实际操作这个过程 .

我的问题归结为:当使用Tensorflow作为Keras的后端时,如何将Keras中的一个模型的预测与多个gpus并行化?

另外,我很好奇是否只有一个gpu可以进行类似的预测并行化 .

高级描述或代码示例将不胜感激!

谢谢!

2 回答

  • 4

    我创建了一个简单的例子来展示如何跨多个gpus运行keras模型 . 基本上,创建了多个进程,每个进程都拥有一个gpu . 要指定进程中的gpu id,设置env变量CUDA_VISIBLE_DEVICES是一种非常简单的方法(os.environ [“CUDA_VISIBLE_DEVICES”]) . 希望这个git repo可以帮到你 .

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

  • 1

    您可以使用此功能来并行化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)
    

相关问题