在上一篇博客【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()
这样之后你会发现模型的大小从GB锐减到几十KB。
