【tensorflow】TF1.x保存.pb模型解决模型越训练越大问题

这篇具有很好参考价值的文章主要介绍了【tensorflow】TF1.x保存.pb模型解决模型越训练越大问题。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

在上一篇博客【tensorflow】TF1.x保存与读取.pb模型写法介绍介绍的保存.pb模型方法中，保存的是模型训练过程中所有的参数，而且训练越久，最终保存的模型就越大。我的模型只有几千参数，可是最终保存的文件有1GB。。。。

但是其实我只想要保存参数去部署模型，然后预测。网上有一些解决方案但都不是我需要的，因为我要用Java部署模型，python这里必须要用builder.add_meta_graph_and_variables来保存参数。以下是解决方案：

举例：模型定义如下

# 定义模型
with tf.name_scope("Model"):
    """MLP"""
    # 13个连续特征数据（13列）
    x = tf.placeholder(tf.float32, [None,13], name='X') 
    # 正则化
    x_norm = tf.layers.batch_normalization(inputs=x)
    # 定义一层Dense
    dense_1 = tf.layers.Dense(64, activation="relu")(x_norm)
    
    """EMBED"""
    # 离散输入
    y = tf.placeholder(tf.int32, [None,2], name='Y')
    # 创建嵌入矩阵变量
    embedding_matrix = tf.Variable(tf.random_uniform([len(vocab_dict) + 1, 8], -1.0, 1.0))
    # 使用tf.nn.embedding_lookup函数获取嵌入向量
    embeddings = tf.nn.embedding_lookup(embedding_matrix, y)
    # 创建 LSTM 层
    lstm_cell = tf.nn.rnn_cell.LSTMCell(64)
    # 初始化 LSTM 单元状态
    initial_state = lstm_cell.zero_state(tf.shape(embeddings)[0], tf.float32)
    # 将输入数据传递给 LSTM 层
    lstm_out, _ = tf.nn.dynamic_rnn(lstm_cell, embeddings, initial_state=initial_state)
    # 定义一层Dense
    dense_2 = tf.layers.Dense(64, activation="relu")(lstm_out[:, -1, :])
    
    """MERGE"""
    combined = tf.concat([dense_1, dense_2], axis = -1)
    pred = tf.layers.Dense(2, activation="relu")(combined)
    pred = tf.layers.Dense(1, activation="linear", name='P')(pred)
    
    z = tf.placeholder(tf.float32, [None, 1], name='Z')

虽然写这么多，但是上面模型的输入只有x、y、z，输出只有pred。所以我们保存、加载模型时，只用考虑这几个变量就可以。

模型保存代码

这里的保存方法建议对比上一篇博客【tensorflow】TF1.x保存与读取.pb模型写法介绍介绍的保存.pb模型方法来看。

import tensorflow as tf
from tensorflow import saved_model as sm
from tensorflow.tools.graph_transforms import TransformGraph
from tensorflow.core.framework import graph_pb2

def get_node_names(name_list, nodes_list):
    name_list.extend([n.name.split(":")[0] for _, n in nodes_list.items() if n.name.split(":")[0] != ''])

# 创建 Saver 对象
saver = tf.train.Saver()

# 生成会话，训练STEPS轮
with tf.Session() as sess:
    # 初始化参数
    sess.run(tf.global_variables_initializer())
    
    ...... # 模型训练逻辑
            
   # 准备存储模型
    path = 'pb_model/'
    # 创建 Saver 对象，用于保存和加载模型的变量
    pb_saver = tf.train.Saver(var_list=None)
    # 将 Saver 对象转换为 SaverDef 对象
    saver_def = pb_saver.as_saver_def()
    # 从会话的图定义中提取包含恢复操作的子图
    saver_def_ingraph = tf.graph_util.extract_sub_graph(sess.graph.as_graph_def(), [saver_def.restore_op_name])
    
    # 构建需要在新会话中恢复的变量的 TensorInfo protobuf
    # 自定义 根据自己的模型来写
    inputs = {
        'x' : sm.utils.build_tensor_info(x),
        'y' : sm.utils.build_tensor_info(y),
        'z' : sm.utils.build_tensor_info(z)
    }
    outputs = {
        'p' : sm.utils.build_tensor_info(pred)
    }

    # 获取节点的名称
    input_node_names = []
    get_node_names(input_node_names, inputs)
    output_node_names = []
    get_node_names(output_node_names, outputs)
    
    # 获取当前会话的图定义
    input_graph_def = sess.graph.as_graph_def()
    
    # 定义需要应用的图转换操作列表
    transforms = ['add_default_attributes',
                  'fold_constants(ignore_errors=true)',
                  'fold_batch_norms',
                  'fold_old_batch_norms',
                  'sort_by_execution_order',
                  'strip_unused_nodes']
    
    # 应用图转换操作，并获取优化后的图定义
    opt_graph_def = TransformGraph(input_graph_def,
                                   input_node_names,
                                   output_node_names,
                                   transforms)
    
    # 创建新的默认图并导入优化后的图定义
    with tf.Graph().as_default() as graph:
    
        all_names = set([node.name for node in opt_graph_def.node])
        saver_def_ingraph_nodes = [node for node in saver_def_ingraph.node if not node.name in all_names]
        merged_graph_def = graph_pb2.GraphDef()
        merged_graph_def.node.extend(opt_graph_def.node)
        merged_graph_def.node.extend(saver_def_ingraph_nodes)

        # 导入合并后的图定义到新的默认图中
        tf.graph_util.import_graph_def(merged_graph_def, name="")

        builder = sm.builder.SavedModelBuilder(path)


        # 将 graph 和变量等信息写入 MetaGraphDef protobuf
        # 这里的 tags 里面的参数和 signature_def_map 字典里面的键都可以是自定义字符串，也可用tf里预设好的方便统一
        builder.add_meta_graph_and_variables(
            sess, tags=[sm.tag_constants.SERVING],
            signature_def_map={sm.signature_constants.PREDICT_METHOD_NAME: SignatureDef},
            saver=pb_saver,
            main_op=tf.local_variables_initializer()
        )

        # 将 MetaGraphDef 写入磁盘
        builder.save()