Tensorflow-在其他模型中重用冻结图

时间:2019-05-29 06:52:46

标签: python tensorflow deeplab

背景:我有一个现有模型,其中正在使用图卷积。除了已有的功能之外,我还想使用语义分割来改善结果。

我想我可以通过下载Deeplab的一些预训练模型来解决此问题,该模型提供了冻结图(.pb文件)以及检查点。我认为冻结的图会更好,因为我不需要复制整个deeplab存储库。

在演示笔记本之后,我做了:

import tensorflow as tf


class DeepLabModel(object):
    INPUT_TENSOR_NAME = 'ImageTensor:0'
    OUTPUT_TENSOR_NAME = 'SemanticPredictions:0'
    MODEL_PATH = "data/deeplab/xception65/frozen_inference_graph.pb"

    def __init__(self):

        self.graph = tf.Graph()

        with open(self.MODEL_PATH, "rb") as file_handle:
            graph_def = tf.GraphDef.FromString(file_handle.read())

        with self.graph.as_default():
            tf.import_graph_def(graph_def, name='')

        self.sess = tf.Session(graph=self.graph)

    def run(self, images):

        return self.sess.run(
            self.OUTPUT_TENSOR_NAME,
            feed_dict={self.INPUT_TENSOR_NAME: images}
        )

现在,这当然不起作用,因为images中的run参数是张量。所以我有两个问题:

  1. 是否可以将张量传递到feed_dict
  2. 有没有更好的方法可以在另一个模型中重用预先训练的模型?

预先感谢

0 个答案:

没有答案