将tf1.x saved_model.pb重新保存到新的tf2.0 saved_model.pb中

时间:2019-12-27 17:46:51

标签: python tensorflow deep-learning protocol-buffers tensorflow2.0

我有一个训练有素的旧tf1.x模型(称为 Model1 ),该模型由占位符,tf.contrib等构建。我可以通过从tf.Session(在tf1.x中)的.ckpt检查点还原图形来使用此模型。 我解决了使用 Model1 的最简单方法是导出它:

# tf1.x code
tf.saved_model.simple_save(sess, saved_Model1_path,
    inputs={'input':'Placeholder:0'}, outputs={'output':'.../Sigmoid:0'})

即使在tf2.0中,我也可以使用获得的save_model.pb:

# tf2.0 code
Model1 = tf.saved_model.load(saved_Model1_path)
out = Model1.signatures['serving_default'](tf.convert_to_tensor(data))['output'].numpy()
out = Model1.signatures['serving_default'].prune('Placeholder:0', '.../Sigmoid:0')(data)
out = Model1.prune('Placeholder:0', '.../Sigmoid:0')(data)

现在想像一下,我有一个用tf2.0 tf.function编写的前/后处理。

我希望将预处理->模型1->后处理的构造导出到tf2.0中的单个save_model.pb中。 这是由于 Model1 的save_model.pb使用tf.Placeholders(有点像这样,我不是专家)引起的问题。

同时,我可以轻松地从其他tf2.0导出的模型中创建saved_model.pb:

import os
import tensorflow as tf
assert tf.__version__[0] == '2'

class M1(tf.Module):
    def __init__(self):
        super(M1, self).__init__()
        self.v = tf.Variable(2.)
    @tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
    def M1_func(self, x):
        return x * self.v

# build some saved_model.pb
m1 = M1()
path_1 = './save1'
path_to_save = os.path.realpath(path_1)
tf.saved_model.save(m1, path_to_save)

# load built saved_model.pb and check it works
m1 = tf.saved_model.load(path_1)
assert 6 == m1.M1_func(3.).numpy()

# build other saved_model.pb using first saved_model.pb as a part of computing graph
class M2(tf.Module):
    def __init__(self):
        super(M2, self).__init__()
        self.run = m1
        self.v = tf.Variable(3.)
    @tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
    def M2_func(self, x):
        return self.run.M1_func(x) * self.v
m2 = M2()
path_2 = './save2'
path_to_save = os.path.realpath(path_2)
tf.saved_model.save(m2, path_to_save)

m2 = tf.saved_model.load(path_2)
assert 18 == m2.M2_func(3.).numpy()

但是,当我尝试执行相同的操作时,除了在tf1.x保存时从tf2.0保存中替换掉首个save_model.pb之外,这是行不通的:

# save first saved_model.pb with tf1.x
import tensorflow as tf
assert tf.__version__[0] == '1'
inp = tf.placeholder(shape=[],dtype=tf.float32)
a = tf.Variable(1.5)
out = a*inp
sess = tf.Session()
sess.run(tf.global_variables_initializer())
assert 7.5 == out.eval({inp:5.}, sess)
path_3 = './save3'
path_to_save = os.path.realpath(path_3)
tf.saved_model.simple_save(sess, path_to_save, inputs={'input': inp}, outputs={'output': out})

现在切换到tf2.0并尝试使用第一个作为计算图的一部分来构建新的saved_model.pb:

import os
import tensorflow as tf
assert tf.__version__[0] == '2'

path_3 = './save3'
path_to_save = os.path.realpath(path_3)
m1 = tf.saved_model.load(path_to_save)

class M2(tf.Module):
    def __init__(self):
        super(M2, self).__init__()
        self.run = m1.signatures['serving_default'].prune('Placeholder:0', 'mul:0')
        self.v = tf.Variable(3.)

    @tf.function(input_signature=[tf.TensorSpec([], tf.float32)])
    def M2_func(self, x):
        return self.run(x) * self.v

m2 = M2()
assert 22.5 == m2.M2_func(5.) # ofc eager execution works

# now save M2 to saved_model.pb and check it works (it does not)
path_4 = './save4'
path_to_save = os.path.realpath(path_4)
tf.saved_model.save(m2, path_to_save)
m2 = tf.saved_model.load(path_4)
m2.M2_func(5.)  # error:
tensorflow.python.framework.errors_impl.FailedPreconditionError:  Attempting to use uninitialized value StatefulPartitionedCall/StatefulPartitionedCall/Variable
     [[{{node StatefulPartitionedCall/StatefulPartitionedCall/Variable/read}}]] [Op:__inference_restored_function_body_207]

Function call stack:
restored_function_body

所以问题是:如何在tf2.0中的单个save_model.pb中保存该体系结构

预处理(tf2.0 @ tf.function)-> Model1 (在tf1.x中创建的saved_model.pb)->后处理 (tf2.0 @ tf.function)

1 个答案:

答案 0 :(得分:0)

问题已解决。查看此导出功能以及如何使用它。此函数实现接受一个输入张量名称和一个张量名称列表。

import tensorflow as tf

def export_tf1(session, in_tnsr_fullname, out_tnsrS_fullname, export_dir='./export'):
    assert isinstance(in_tnsr_fullname, str)
    assert all([isinstance(out_tnsr_fullname, str) for out_tnsr_fullname in out_tnsrS_fullname])

    in_tnsr_name = in_tnsr_fullname.split(':')[0]
    out_tnsrS_name = [out_tnsr_fullname.split(':')[0] for out_tnsr_fullname in out_tnsrS_fullname]

    graph_def = tf.graph_util.convert_variables_to_constants(session, session.graph.as_graph_def(), out_tnsrS_name)

    tf.reset_default_graph()
    outs = tf.import_graph_def(graph_def, name="", return_elements=out_tnsrS_fullname)
    g = outs[0].graph

    builder = tf.saved_model.builder.SavedModelBuilder(export_dir)

    with tf.Session(graph=g) as sess:
        input_signatures = {in_tnsr_name: g.get_tensor_by_name(in_tnsr_fullname)}
        output_signatures = {}
        for out_tnsr_name, out_tnsr_fullname in zip(out_tnsrS_name, out_tnsrS_fullname):
            output_signatures[out_tnsr_name] = g.get_tensor_by_name(out_tnsr_fullname)
        signature = tf.saved_model.signature_def_utils.predict_signature_def(input_signatures, output_signatures)

        builder.add_meta_graph_and_variables(
            sess,
            [tf.saved_model.tag_constants.SERVING],
            {tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature},
            clear_devices=True
        )

    builder.save()

如何使用导出功能从tf1_ckpt检查点接收.pb:

import tensorflow as tf
assert tf.__version__[0] == '1'

g = tf.get_default_graph()
sess = tf.Session(graph=g)
ckpt_tf1_path = 'some_directory/name.ckpt'  # just an example
tf.train.Saver().restore(sess, ckpt_tf1_path)
input_tensor_name = 'x_tnsr:0'  # just an example
out_tensor_name = 'y_tnsr:0'  # just an example
export_tf1(sess, input_tensor_name, [out_tensor_name], export_dir)

如何在使用tf2.0的.pb中重用tf1_ckpt中的.pb:

import tensorflow as tf
assert tf.__version__[0] == '2'

class Export(tf.Module):
    def __init__(self):
        super(Export, self).__init__()
        tf1_saved_model_directory = 'directory/saved_model'  # just an example
        self.tf1_model = tf.saved_model.load(tf1_saved_model_directory)
        input_tensor_name = 'x_tnsr:0'  # just an example
        out_tensor_name = 'y_tnsr:0'  # just an example
        self.tf1_model = self.tf1_model.prune(input_tensor_name, out_tensor_name)

    @tf.function
    def __call__(self, x):
        out = self.tf1_model(x)
        return out

export_dir = './saved_model'
tf.saved_model.save(Export(), export_dir)