将模型从 tensorflow SavedModel 加载到多个 GPU

Question

假设有人递给我一个 TF SavedModel，我想在我机器上的 4 个 GPU 上复制这个模型，这样我就可以运行对批量数据进行并行推理。有没有很好的例子说明如何做到这一点？

我可以通过这种方式加载保存的模型：

def load_model(self, saved_model_dirpath):
    '''Loads a model from a saved model directory - this should 
       contain a .pb file and a variables directory'''

    signature_key = tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
    input_key = 'input'
    output_key = 'output'

    meta_graph_def = tf.saved_model.loader.load(self.sess, [tf.saved_model.tag_constants.SERVING],
                                                saved_model_dirpath)
    signature = meta_graph_def.signature_def

    input_tensor_name = signature[signature_key].inputs[input_key].name
    output_tensor_name = signature[signature_key].outputs[output_key].name

    self.input_tensor = self.sess.graph.get_tensor_by_name(input_tensor_name)
    self.output_tensor = self.sess.graph.get_tensor_by_name(output_tensor_name)

..但这需要我有会话句柄。对于我自己编写的模型，我可以访问推理函数，我可以调用它并使用 with tf.device() 包装它，但在这种情况下，我不确定如何从中提取推理函数保存的模型。我应该加载 4 个单独的会话还是有更好的方法？找不到太多关于此的文档，但如果我遗漏了什么，请提前致歉。谢谢！

Answer 1

TensorFlow 目前不支持此用例。不幸的是，"replicating the inference function" 仅基于 SavedModel（它基本上是具有一些元数据的计算图）是一个相当复杂（如果实施的话也很脆弱）的图转换问题。

如果您无法访问生成此模型的源代码，最好的办法是将 SavedModel 加载到 4 个单独的图形中 4 次，每次将目标设备重写到相应的 GPU。然后，运行每个graph/session分开。

请注意，您可以同时调用 sess.run() 多次，因为 sess.run() 会在实际计算时释放 GIL。您只需要几个 Python 线程。

将模型从 tensorflow SavedModel 加载到多个 GPU

Loading a model from tensorflow SavedModel onto mutliple GPUs

multi-gpu

tensorflow