tf.estimator.export_saved_model保存模型

之前使用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，你也可以输出最终概率值，根据自己需要定义模型输入输出即可。