之前使用tensorflow框架做学术研究时多是生成ckpt文件,但是实际工作中,模型上线部署常用 SavedModel 保存模型,生成的模型是封闭的序列化格式,它具有语言独立性,可独立运行,并且可以简单被其他语言(python,java等)调用,SavedModel 格式是tf 2.0 推荐的格式,支持tf-serving部署,部署时很方便。好了介绍了这么多,就是希望告诉大家使用tf框架时最好直接存储成PB文件,方便之后使用。
SavedModel
我们首先看下SavedModel中都包含了什么
saved_model.pb:保存的是网络图结构,同ckpt文件夹中的*.meta文件
variable:参数权重,包含了模型的所有变量
保存模型
以微调bert预训练模型为例:
do_export = True
def serving_input_fn():
label_ids = tf.placeholder(tf.int32, [None], name='label_ids')
input_ids = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_ids')
input_mask = tf.placeholder(tf.int32, [None, FLAGS.max_seq_length], name='input_mask')
segment_ids = tf.placeholder(tf.int32, [FLAGS.max_seq_length], name='segment_ids')
input_fn = tf.estimator.export.build_raw_serving_input_receiver_fn({
'label_ids': label_ids,
'input_ids': input_ids,
'input_mask': input_mask,
'segment_ids': segment_ids,
})()
return input_fn
if do_export:
estimator._export_to_tpu = False
estimator.export_savedmodel(output_dir, serving_input_fn)
加载模型
import tensorflow as tf
predict_fn = tf.contrib.predictor.from_saved_model(MODEL_DIR)
#feature.*模型的输入
prediction = predict_fn({
"input_ids": feature.input_ids,
"input_mask": feature.input_mask,
"segment_ids": feature.segment_ids,
})
# print(prediction)
eval_logits = prediction["logits"]#输出
在这里,我输出的是模型最终的logits,你也可以输出最终概率值,根据自己需要定义模型输入输出即可。