将tensorflow SavedModel中的模型加载到多个GPU

假设有人递给我一个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)

---

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

目前TensorFlow中不支持此用例。不幸的是，“复制推理函数”仅基于SavedModel（基本上是带有一些元数据的计算图），这是一个相当复杂的（如果实现的话，也是脆弱的）图转换问题

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

请注意，您可以同时调用

sess.run（）

多次，因为

sess.run（）

会在实际计算时释放GIL。您只需要几个Python线程