由于TF里面的概念比较接地气，所以用tf1.x保存.pb模型时总是怕有什么操作漏掉了，会造成保存的模型是缺少变量数据或者没有保存图，所以先明确一下：用TF1.x保存模型时只需要保存模型的输入输出的变量（多输入就保存多个），不需要保存中间的变量；用TF1.x加载模型时只需要加载保存的模型，然后读一下输入输出变量（多输入就读多个），不需要初始化（反而会重置掉变量的值）。
举例：模型定义如下
# 定义模型 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。所以我们保存、加载模型时，只用考虑这几个变量就可以。
模型保存代码
import tensorflow as tf from tensorflow import saved_model as sm # 创建 Saver 对象 saver = tf.train.Saver() # 生成会话，训练STEPS轮 with tf.Session() as sess: # 初始化参数 sess.run(tf.global_variables_initializer()) ...... # 模型训练逻辑 # 准备存储模型 path = 'pb_model/' dense_model_var = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) pb_saver = tf.train.Saver(dense_model_var) builder = sm.builder.SavedModelBuilder(path) # 构建需要在新会话中恢复的变量的 TensorInfo protobuf # 自定义 根据自己的模型来写 X = sm.utils.build_tensor_info(x) Y = sm.utils.build_tensor_info(y) Z = sm.utils.build_tensor_info(z) P = sm.utils.build_tensor_info(pred) # 构建 SignatureDef protobuf # inputs outputs 自定义 根据自己的模型来写 SignatureDef = sm.signature_def_utils.build_signature_def( inputs={'X': X, 'Y': Y, 'Z': Z}, # 可用sm.signature_constants.PREDICT_INPUTS outputs={'P': P}, # 可用sm.signature_constants.PREDICT_OUTPUTS method_name="tensorflow/serving/predict" ) # 将 graph 和变量等信息写入 MetaGraphDef protobuf # 这里的 tags 里面的参数和 signature_def_map 字典里面的键都可以是自定义字符串，也可用tf里预设好的方便统一 builder.add_meta_graph_and_variables(sess, tags=['serve'], signature_def_map={ sm.signature_constants.PREDICT_METHOD_NAME: SignatureDef}, saver=pb_saver, main_op=tf.local_variables_initializer()) # 将 MetaGraphDef 写入磁盘 builder.save()
最重要的是这一句：dense_model_var = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES)，意思是保存当前作用域下的所有可训练的变量。
我之前写的是dense_model_var = tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES, name_scope="Model")，这样读不了所有的可训练变量，只能读到embedding_matrix 一个，虽然也能保存模型，但是没保存模型的其他变量值，就会出错。
模型加载代码
import tensorflow as tf from tensorflow import saved_model as sm tf.reset_default_graph() # 创建一个新的默认图 graph = tf.Graph() # 需要建立一个会话对象，将模型恢复到其中 with tf.Session(graph=graph) as sess: path = 'pb_model/' MetaGraphDef = sm.loader.load(sess, tags=['serve'], export_dir=path) # 解析得到 SignatureDef protobuf SignatureDef_map = MetaGraphDef.signature_def SignatureDef = SignatureDef_map[sm.signature_constants.PREDICT_METHOD_NAME] # 解析得到 3 个变量对应的 TensorInfo protobuf X = SignatureDef.inputs['X'] Y = SignatureDef.inputs['Y'] Z = SignatureDef.inputs['Z'] P = SignatureDef.outputs['P'] # 解析得到具体 Tensor # .get_tensor_from_tensor_info() 函数中可以不传入 graph 参数，TensorFlow 自动使用默认图 # x = sm.utils.get_tensor_from_tensor_info(X) # y = sm.utils.get_tensor_from_tensor_info(Y) # z = sm.utils.get_tensor_from_tensor_info(Z) x = sess.graph.get_tensor_by_name(X.name) y = sess.graph.get_tensor_by_name(Y.name) z = sess.graph.get_tensor_by_name(Z.name) p = sess.graph.get_tensor_by_name(P.name) # 这里就可以开始进行预测或者继续训练了 TODO total_loss = sess.run(loss_function, feed_dict={x: dense_ch_val, y: sparse_ch_val, z: score_val}) print(total_loss)
