将.pb从Tensorflow 1.14迁移到2.0

时间:2019-10-17 04:51:15

标签: python tensorflow tensorflow2.0

我正在尝试在Tensorflow 2.0中运行FaceNet模型。我已经下载了一组预训练的权重(.pb文件),并用于使用tf.GraphDef()在TF 1.14中加载图形。

我想知道我应该如何在TF 2.0中运行它:

我尝试使用tf.SavedModel.load()方法,但这会返回一个空的签名字典。

是否可以在新的Tensorflow 2.0版本中重用TF 1.x pb文件?如果可以的话?

1 个答案:

答案 0 :(得分:0)

好吧,根据这样的说法:https://github.com/tensorflow/community/blob/master/sigs/testing/faq.md看来

“从TensorFlow 1.x模型WONT生成的冻结图可在TF 2.0中工作”

我确实找到了一种将模型转换为保存的模型的方法:

class Net(object):

    def __init__(self, model_path):
        self.model_filepath = model_path
        self.load_graph(model_filepath=self.model_filepath)


    def load_graph(self, model_filepath):
        print("Loading model...")


        with tf.gfile.GFile(model_filepath, 'rb') as f:
            graph_def = tf.GraphDef()
            graph_def.ParseFromString(f.read())

        with tf.Session() as sess:
            with tf.Graph().as_default() as graph:
                tf.import_graph_def(graph_def, name='')
                signature = tf.saved_model.signature_def_utils.predict_signature_def(
                    inputs={'image_batch':graph.get_tensor_by_name('image_batch:0'),
                            'phase_train': graph.get_tensor_by_name('phase_train:0')},
                    outputs={'embeddings': graph.get_tensor_by_name('embeddings:0')}
                )
                builder = tf.saved_model.builder.SavedModelBuilder("./Output/")
                builder.add_meta_graph_and_variables(
                    sess=sess,
                    tags=[tf.saved_model.tag_constants.SERVING],
                    signature_def_map={'serving_default': signature}
                )
                builder.save()


Net(path_to_pb_file)