在上一篇博客【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。