tensorflow模型保存（三）——tensorflow1.x版本的savedmodel格式的模型保存与加载

前言：

Tensorflow的保存分为四种：

• 1. checkpoint模式；
• 2. saved_model模式（包含pb文件和variables）；
• 3. 纯pb模式；（只有一个pb文件）
• 4.keras的 h5 模式

前面的两篇文章已经分别介绍了高层API与低层API关于各种模型的保存，比如在tensorflow2.x中如何保存checkpoint与savedmodel，在tensorflow1.x中如何保存checkpoint，本文着重介绍在tensorflow1.x版本中怎么保存savedmodel。

前面的两篇文章请参考：

详解tensorflow2.0的模型保存方法（一）

tensorflow中的检查点checkpoint详解（二）——以tensorflow1.x 的模型保存与恢复为主

本文介绍在tensorflow1.x版本中如何使用savedmodel格式的模型，本文基于tensorflow1.13.1

 

一、tensorflow1.x版本的savedmodel模型保存与恢复

1.1 SavedModel与ckpt在使用中的共同点

Tensorflow训练的模型可以保存为ckpt格式，但是这种格式的模型文件在跨语言方面不是很灵活。而且在做模型的restore然后进行预测的时候，必须要知道两个tensor_name，即：

最核心的就在于要给模型的输入与输出起一个易于辨识的名称，这样词能够很好的恢复模型并进行预测测试。

实际上，在tensorflow1.x的版本中，这一点与ckpt的恢复使用是一致的，我必须要能够明确地知道整个graph的输入与输出的那一个operation的名称才行，在前面一篇文章，即

tensorflow中的检查点checkpoint详解（二）——以tensorflow1.x 的模型保存与恢复为主

中，模型的输入与输出是：

输入：model_input
输出：dense4_output

1.2 SavedModel的结构

以SavedModel格式保存模型时，tensorflow将创建一个SavedModel目录，该目录由以下子目录和文件组成：

assets/
assets.extra/
variables/
    variables.data-?????-of-?????
    variables.index
saved_model.pb

其中，各目录和文件的说明如下：

• assets/是包含辅助（外部）文件（如词汇表）的子文件夹。资产被复制到SavedModel位置，并且可以在加载特定的MetaGraphDef时读取。
• assets.extra是一个子文件夹，高级库和用户可以将自己的资源添加进去，这些资源将与模型共存但不由图形加载。此子文件夹不由SavedModel库管理。
• variables/是包含tf.train.saver输出的子文件夹。
• saved_model.pb或saved_model.pbtxt是SavedModel协议缓冲区。它将图形定义作为MetaGraphDef协议缓冲区。

assets/和assets.extra目录是可选的。

checkpoint与savedmodel其实也很结构上面相似的地方，

• 第一：在checkpoint中，graph的结构是存储在 xxx.meta文件中的，而真正的变量参数是存储的data文件中；savedmodel其实也是类似的， xxx.pb 文件仅仅保存graph的结构，真正的变量是存储的variables里面的，
• 第二：但是也有不同的地方，  xxx.pb  文件还可以保存常量，所以我可以将变量的值转化为常量，也就是所谓的freezing，然后整个graph结构以及各个常量、变量的值都可以存储在一个 pb 文件中了；
• 第三：另外，pb文件还可以使用签名，方便graph的重构与恢复

1.3 tensorflow1.x中SavedModel模型的保存

比如我有下面的模型代码：

X = tf.placeholder(dtype=tf.float32,shape=(None,210),name="model_input")
Y = tf.placeholder(dtype=tf.float32,shape=(None,10),name="model_output")
model = KeypointNetwork()        # 构建分类网络，这个是自己编写通过面向对象的方式实现的网络结构
Y_pred = model(X)       # 得到模型的输出，输出的是分类之前的logits，在graph中的名称是dense4_output

怎么保存呢？

实际上就两句话的事儿

with tf.Session(config = config) as sess:
     for epoch in range(epochs): 
         # 训练代码
     tf.saved_model.simple_save(sess,"./saved_model",inputs={"model_input":X },outputs={"dense4_output": Y_pred})
     print("model has saved,model format is saved_model !")

（1）保存方式一——tf.saved_model.simple_save()

函数原型如下：

simple_save(session,      # 要保存的会话
            export_dir,   # saved_model 所存储的文件夹
            inputs={"x": x, "y": y},  # 整个graph的输入，输入需要与定义的输入operation匹配哦
            outputs={"z": z})  # 整个graph的输出，输入需要与定义的operation匹配哦

注意：

• （1）在graph中，输入输出的名称必须要保持一致。在上面的例子中，我的输入是X，输出是Y_pred，但是在定义的时候他们的operation的名称分别是model_input和dense_output, 所以在保存的时候输入和输出的名称也必须用这两个名称，保持统一；
• （2）文件夹不能事先存在。需要保存的文件夹，即./saved_model文件夹不能事先存在，否则会提示已经存在了，不能保存模型。

保存结束之后，会得到下面的模型文件

|---------saved_model

          |---------variables

                     |---------variables.data-00000-of-00001

                     |---------variables.index

          |--------saved_model.pb

 

(2) tensorflow1.x第二种保存saved_model的方法

import tensorflow as tf
from tensorflow.saved_model.signature_def_utils import predict_signature_def
from tensorflow.saved_model import tag_constants


# tensorflow1.x中另一种保存saved_model的方法
builder = tf.saved_model.builder.SavedModelBuilder("./saved_model")
signature = predict_signature_def(inputs={'Input': x},
                                  outputs={'Output': y})
builder.add_meta_graph_and_variables(sess=sess,
                                     tags=[tag_constants.SERVING],
                                     signature_def_map={'predict': signature})
builder.save()

函数原型如下：

add_meta_graph_and_variables(
    sess,     # session会话对象
    tags,     # tensorflow.saved_model.tag_constants 中所定义的标签
    signature_def_map=None, 
    assets_collection=None, 
    legacy_init_op=None,
    clear_devices=False,
    main_op=None, 
    strip_default_attrs=False, 
    saver=None
)

需要注意的几个地方，遵循一个四步走策略：

（1）第一种方法

它实际上是一种签名，即我需要先声明这个graph的输入、输出是什么？常见的定义方式为：

from tensorflow.saved_model import signature_constants, signature_def_utils, tag_constants, utils

# 第一步：创建签名，模型签名的定义
model_signature = signature_def_utils.build_signature_def(
        inputs = {"input": utils.build_tensor_info(tensor)},   # 输入的那个tensor,输出的那个tensor，比如：y=tf.matmul(...)
        outputs = {"output": utils.build_tensor_info(tensor)}, # 这里的input和output表示的是输入输出的operation_name

# 第二步：创建signature_def_map的 key
method_name=signature_constants.PREDICT_METHOD_NAME)

# 第三步创建builder   
builder = saved_model_builder.SavedModelBuilder(export_path)

# 第四步：将meta和variables添加进builder里面
builder.add_meta_graph_and_variables(
            session, 
            [tag_constants.SERVING],
            clear_devices=True,
            signature_def_map = {
                    signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY  model_signature
                    }
            )

（2）第二种方法

import tensorflow as tf
from tensorflow.saved_model.signature_def_utils import predict_signature_def
from tensorflow.saved_model import tag_constants


# tensorflow1.x中另一种保存saved_model的方法
builder = tf.saved_model.builder.SavedModelBuilder("./saved_model")

signature = predict_signature_def(inputs={'Input': x},   # 这里的x,y指的都是某一个operation哦
                                  outputs={'Output': y}) # 所以这里的Input、Output实际上是operation_name
builder.add_meta_graph_and_variables(sess=sess,
                                     tags=[tag_constants.SERVING],
                                     signature_def_map={'predict': signature})
builder.save()

注意：上面的两种方式实际上是一种，只不过有轻微的区别，还有一个地方就是

（1）对于参数 tags 

它是tag_constant 中的一个常量来指定的，里面包含了四个常量，分别是

GPU = 'gpu'
SERVING = 'serve'
TPU = 'tpu'
TRAINING = 'train'

（2）对于上面的起的几个名称 Input、Output、predict 这几个名称是自己取的，不再需要和graph中的tensor_name一样，如果还是要保持一样的tensor_name，那么这个signature岂不是毫无意义了。实际上这个关键的signature_def_map的形式就是像下面这个样子的：

{'predict': inputs {
  key: "Input"
  value {
    name: "model_input:0"
    dtype: DT_FLOAT
    tensor_shape {
      dim {
        size: -1
      }
      dim {
        size: 210
      }
    }
  }
}
outputs {
  key: "Output"
  value {
    name: "dense4_output:0"
    dtype: DT_FLOAT
    tensor_shape {
      dim {
        size: -1
      }
      dim {
        size: 10
      }
    }
  }
}
method_name: "tensorflow/serving/predict"
}

我们发现，这对于我们后面实用签名signature进行回复是有作用的，它里面存储了整个graph的签名信息，又包含了输入张量名称、输出张量名称，这是我们恢复模型的依据

怎么解析signature呢？

第一：signature总的名称是“predict”，可以通过字典的方式去获取它的值，它的值包含有两个属性，一个属性是inputs，表示多个输入，上面的例子中只有一个输入，然后inputs里面又是一个字典，对应不同的几组key-value值，这个key-value就是我们创建签名的时候所构造的，key的名称实际上表示的是一个operation_name，最后通过value的name属性，可以获取该输入所对应的tensor_name名称，这就解析出来了，对于多个输出outputs也是同样的格式和道理。

比如我们常常这样来获取输入与输出的operation_name与tensor_name：

# 获取输入的tensor_name
input_tensor_name = meta_graph.signature_def["predict"].inputs["Input"].name
# 获取输出的tensor_name
output_tensor_name = meta_graph.signature_def["predict"].outputs["Output"].name

但是这样相当于我还是要知道我在创建签名signature的时候指定的signature的名称及predict以及自己定义输入输出的operation_name，所以并没有比直接获得tensor_name来得方便，如果我根本就不知道签名是怎么定义的呢？那该怎么恢复？所以我们更常用的是采用下面的办法，不通过名称，直接通过索引来取得名称，不用关心我在定义签名的时候到底取了一些什么名字，这才是真正的“名称解耦”：

# 加载具有signature的pb文件
meta_graph = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], model_file_path)

# meta_graph.signature_def 是一个字典，通过item()转化成 [(key,value)]的列表
# [0] 表示的获取整个signature的内容，即等价于上面的（"predict",inputs...outpots...）
# [1] 表示获取签名的值，即同时包值inputs属性和outputs属性
inputs_outputs_signature = list(meta_graph.signature_def.items())[0][1]

# 存储输出的 tensor_name 和 operation_name
output_tensor_names = [] 
output_op_names = []

# 遍历输出属性outputs，获取输出的operation_name和tensor_name,outputs中是多个key-value对
for output_item in inputs_outputs_signature.outputs.items():
    output_op_name = output_item[0]          # 即自己定义的时候取的输出的名称Output
    output_op_names.append(output_op_name)
            
    output_tensor_name = output_item[1].name # 即输出的那个tensor_name，即 dense4_output
    output_tensor_names.append(output_tensor_name)

inputs_feed_dict = {}  # 构造输入的feed_dict
for input_item in inputs_outputs_signature.inputs.items():
    input_op_name = input_item[0]
    input_tensor_name = input_item[1].name
    feed_dict_map[input_tensor_name] = test_x

predict_y = sess.run(output_tensor_names, feed_dict=inputs_feed_dict)
            

本质上其实就是字典的一步一步解析，当然我也可以有其他的灵活的解析方式。

 

1.4 tensorflow1.x中saved_model模型的恢复与预测

（1）通过指定输入、输出的tensor_name来实现

重点：同ckpt一样，我需要整个graph的输入与输出的那个operation名称，而这个名称是在自己定义graph的时候指定的。

代码如下：

tf.saved_model.loader.load(sess, ["serve"], "./saved_model")
    graph = tf.get_default_graph()
    # 获取输入输出的tensor_name，而这个是在定义自己的operation的时候指定的
    x = sess.graph.get_tensor_by_name('model_input:0')
    y = sess.graph.get_tensor_by_name('dense4_output:0')
            
    result = sess.run(y,feed_dict={x: test_x})

函数接口如下：

tf.compat.v1.saved_model.load(
    sess,    # 会话对象
    tags,    # tensorflow.saved_model.tag_constants 中所定义的标签
    export_dir, # saved_model所保存的文件夹
    import_scope=None, 
    **saver_kwargs
)

上面的这种解析方法适用于以上两种保存方法保存的savemodel格式的模型，但是有一个问题，那就是需要指定输入和输出的tensor_name，有没有什么办法不需要呢？实际上上面的第二种保存方法里面的签名signature就提供了解决方案，

那么如何可以在不知道tensor name的情况下使用呢，实现彻底的解耦呢？ 给add_meta_graph_and_variables方法传入第三个参数，signature_def_map即可，那具体怎么恢复呢？

（2）通过保存时指定的signature来恢复savedmodel模型

第一步：获取pb文件中的签名

with tf.Session() as sess:
    model_file_path = "./saved_model_build"
    meta_graph = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], model_file_path)
    # 打印查看签名signature的内容，内容如上面所显示的那样
    print(meta_graph.signature_def)

第二步：根据signature来恢复处输入与输出的tensor_name和operation_name

有下面的一些方法

with tf.Session() as sess:
    # 加载具有signature的pb文件
    meta_graph = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING],     
    model_file_path)

    # meta_graph.signature_def 是一个字典，通过item()转化成 [(key,value)]的列表
    # [0] 表示的获取整个signature的内容，即等价于上面的（"predict",inputs...outpots...）
    # [1] 表示获取签名的值，即同时包值inputs属性和outputs属性
    inputs_outputs_signature = list(meta_graph.signature_def.items())[0][1]

    # 存储输出的 tensor_name 和 operation_name
    output_tensor_names = [] 
    output_op_names = []

    # 遍历输出属性outputs，获取输出的operation_name和tensor_name,outputs中是多个key-value对
    for output_item in inputs_outputs_signature.outputs.items():
        output_op_name = output_item[0]          # 即自己定义的时候取的输出的名称Output
        output_op_names.append(output_op_name)
            
        output_tensor_name = output_item[1].name # 即输出的那个tensor_name，即 dense4_output
        output_tensor_names.append(output_tensor_name)
    
    sess.run(tf.global_variables_initializer())
    sess.run(test_iterator.initializer)
            
    for k in range(test_batch_count):
        test_x,test_y = sess.run([x_test,y_test])
        feed_dict_map = {}
        for input_item in model_graph_signature.inputs.items():
            input_op_name = input_item[0]
             input_tensor_name = input_item[1].name
             feed_dict_map[input_tensor_name] = test_x
        predict_y = sess.run(output_tensor_names, feed_dict=feed_dict_map)

 

二、tf.train.write_graph()方法直接将graph保存成pb文件

前面介绍的两种方法保存成savedmodel格式之后都会得到一个  xxx.pb 文件，一个 variables 文件夹，里面包含数据信息，那有没有办法将所有的数据，包括graph结构与变量全部都存储进一个 xxx.pb 文件中呢？

2.1 直接保存一个 pb 文件

我们也可以使用另外的一种方法来将模型保存为pb文件，如下代码：

with tf.Session() as sess: 
    # 训练代码
    tf.train.write_graph(sess.graph_def, "./new_pb_model", "my_test.pb", False)
    print("pb文件已经保存完毕！")

我们可以发现在new_pb_model文件夹下面出了一个单独的  my_test.pb 文件，只有这一个，但是这文件比较小。函数原型如下：

# 在tensorflow1.13以及之前的版本
tf.train.write_graph(graph_or_graph_def, logdir, name, as_text=True)

# 在之后的版本
tf.io.write_graph(graph_or_graph_def, logdir, name, as_text=True)

但是有一个问题：

tf.train.write_graph将导出模型的GraphDef文件，实际上保存了训练的神经网络的结构图信息。存储格式为protobuffer，所以文件名后缀为pb。但是pb文件：保存图模型的计算流程图，包含常量，不包含变量。

没错，Graphdef中不保存任何 Variable 的信息，所以如果我们从 graph_def 来构建图并恢复训练的话，是不能成功的

既然没有变量信息，自然我们虽然能够恢复出graph的结构，但是却不能够会付出相对应的变量值，那肯定是不行的，那有什么办法呢？既然pb文件可以存储常量constant，可不可以将variables转化成constant呢？这就是解决思路。

2.2 将所有的variable转化成constant保存到一个 pb 文件中——可以实现finetune

由于tensorflow没有提供类似于加载.ckpt文件的restore接口，不能够恢复 pb 文件中的权重信息，况且 pb 文件根本就没有保存variable，所以在做.pb文件用于模型finetune时，需要将模型中的trainable variables全部保存下来，并且在加载.pb文件时需要根据变量名称将变量值一一赋值到模型中。 应用步骤分为保存trainable variables到.pb文件和从.pb文件加载trainable variables：

（1）将constant转化成variables

函数原型：

tf.graph.util.convert_variables_to_constants(
    sess,  
    input_graph_def, 
    output_node_names,
    variable_names_whitelist=None, 
    variable_names_blacklist=None)

参数详解：

• sess: 当前使用的会话对象sess
• input_graph_def: 是一个GraphDef object ，及当前会话中的Graph
• output_node_names: graph输出节点的名称，形如 [“name1”，“name2” ]
• variable_names_whitelist: 需要转化的变量Variable所组成的list，默认情况下graph中的所有variable均会转化成constant(by default, all variables are converted).
• variable_names_blacklist: 忽略转化。即不需要转化成constant的variables所组成的list

一般的使用方法如下：

with tf.Session(config = config) as sess:
    # 其他代码
    var_list = tf.trainable_variables()
    frozen_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, ['dense4_output'],var_list)
    # 写入到 pb文件
    tf.io.write_graph(frozen_graph_def,"pb_model","freeze_eval_graph.pb",as_text=False)
    
    # 写入 pb 文件等价于下面的代码
    with tf.gfile.FastGFile("./freeze_eval_graph.pb", mode='wb') as f:
        f.write(frozen_graph_def.SerializeToString())

（2）从ckpt 模型转化成单个的 pb 文件

with tf.Session(config = config) as sess:
    # 构建模型结构   
    eval_model = TrainKeypointNetwork(is_train=False)
    saver = tf.train.Saver()  # 导入变量参数
    saver.restore(sess, './quant_ckpt_model/quant_keypoint_model.ckpt-9')
    
    # 将变量转化成constant
    frozen_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, ['dense4_output'])
    tf.io.write_graph(frozen_graph_def,"pb_model","freeze_eval_graph.pb",as_text=False)
    print("model have been frozen... ...")

或者也就可以向下面这样

with tf.Session(config = config) as sess:
    # 构建模型结构   
    saver = tf.train.import_meta_graph('./ckpt_model/keypoint_model.ckpt-99.meta')
    # 载入模型参数
    saver.restore(sess,'./ckpt_model/keypoint_model.ckpt-99')
    
    # 将变量转化成constant
    frozen_graph_def = tf.graph_util.convert_variables_to_constants(sess, sess.graph_def, ['dense4_output'])
    tf.io.write_graph(frozen_graph_def,"pb_model","freeze_eval_graph.pb",as_text=False)
    print("model have been frozen... ...")

 

2.3 如何加载整个 pb文件——tf.import_graph_def

with tf.Session(config=config) as sess:
    sess.run(tf.global_variables_initializer())
    sess.run(test_iterator.initializer)
        
    with tf.gfile.FastGFile("pb_model/freeze_eval_graph.pb", 'rb') as f:
        # 使用tf.GraphDef()定义一个空的Graph
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        sess.graph.as_default()
        tf.import_graph_def(graph_def, name='')

        # 获取输入输出的tensor_name
        input_x = sess.graph.get_tensor_by_name('model_input:0')
        pred_y = sess.graph.get_tensor_by_name('dense4_output:0')
            
        result=sess.run(pred_y,feed_dict={input_x:test_x})  #需要的就是模型预测值model_Y，这里存为result

tf.import_graph_def函数原型

# 旧版本
# tf.import_graph_def
# 新版本
tf.graph_util.import_graph_def(
    graph_def, 
    input_map=None, 
    return_elements=None, 
    name=None, 
    op_dict=None,
    producer_op_list=None
)

该函数的返回值默认是None，但是当我们指定参数return_elements之后就可以返回对应的tensor了，因为我们需要graph的输入与输出，所以，我们可以这样做：

with tf.Session(config=config) as sess:
    sess.run(tf.global_variables_initializer())
    sess.run(test_iterator.initializer)
        
    with tf.gfile.FastGFile("pb_model/freeze_eval_graph.pb", 'rb') as f:
        # 使用tf.GraphDef()定义一个空的Graph
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        sess.graph.as_default()
        input_x,pred_y = tf.import_graph_def(graph_def, return_elements=["model_input:0","dense4_output:0"], name='')

        # 获取输入输出的tensor_name,因为上面已经通过return_element返回了需要的输入与输出，这两句话可以不要了
        # input_x = sess.graph.get_tensor_by_name('model_input:0')
        # pred_y = sess.graph.get_tensor_by_name('dense4_output:0')
            
        result=sess.run(pred_y,feed_dict={input_x:test_x})  #需要的就是模型预测值model_Y，这里存为result

但是上面有一个问题，因为上面的 pb 文件中所有的variables均转化成了constant，我只用模型进行预测当然没有问题，因为只涉及到前向运算，不用方向传播，但是如果我想继续之前的训练该怎么办呢？即所谓的finetune，常量是不变的，是没有办法改变它的值的，自然也没办法更新梯度了，

我们可以通过实验查看：

def finetune_pb():
    pb_para_dic = {} # 用于存储pb文件变量信息的字典结构 key-变量名,value-变量值
    with tf.Graph().as_default():
        output_graph_def = tf.GraphDef()
    with open("./pb_model/freeze_eval_graph.pb", "rb") as f:
        output_graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(output_graph_def, name="")  # 将pb文件中保存的graph加载到新创建的graph def中
        # 查看从pb文件中加载得到的graph是否含有变量,发现返回的是空的
        var_list = tf.global_variables()
    
    with tf.Session() as sess:
        constant_values = {}
        constant_ops = [op for op in sess.graph.get_operations() if op.type == "Const"]
        for constant_op in constant_ops:
            print(constant_op.name)

我们可以像下面这样做。

参考：https://zhuanlan.zhihu.com/p/47649285

因为这个地方代码补全，我暂时还没找到合适的解决方案，望有大神告知如何从 单独的  pb 文件恢复变量，然后进行fine tune。

2.4 从.pb文件加载trainable variables用于继续训练（finetune）

 

三、关于pb文件的总结

• （1）本质上依然是存储graph的结构的，并不存储任何变量信息；
• （2）但是pb文件可以存储常量，我们可以将变量转化成常量constant，即所谓的freezing，然后将整个模型结构，模型权重全部存储成一个pb文件
• （3）pb文件可以定值签名signature，方便名称解耦，在恢复模型的时候就不用通过tensor_name了，而是通过签名来实现模型恢复。

另外上面的内容总结如下：

• （1）两种保存savedmodel格式的方法： tf.saved_model.simple_save（）和 tf.saved_model.builder.SavedModelBuilder（）
• 得到的是 xxx.pb 模型文件和 variables变量文件。
• （2）恢复savedmodel文件的方法： meta_graph = tf.saved_model.loader.load（），然后可以通过tensor_name或者是通过signature来去的输入输出，然后进行预测。
• （3）创建单个 pb 文件的方法：tf.train.write_graph（），其中包括了只保存最基本的pb文件、将variable转化成constant再保存、从ckpt来转化成pb文件  三种方法。
• （4）读取 单个的 pb 文件 ，