【云+社区年度征文】tensorflow 2.0 Estimator Keras读取saved model并预测

原创

大鹅

发布于 2020-12-07 17:35:03

8260

文章被收录于专栏：大鹅专栏：大数据到机器学习大鹅专栏：大数据到机器学习

背景

使用tensorflow2.0以上版本框架用Keras或者Estimator方式保存模型有两种方式加载模型并预测。

Keras框架保存模型后可以直接加载并调用predict方法预测；

estimator将比较麻烦，需要签名并传入tensor才可以预测；

Keras模型预测

import tensorflow as tf
from tensorflow import keras
model = tf.keras.models.load_model(export_dir)

# dataframe 特征读取与处理
X = dict(dataframe)
c = model.predict(X)
output = np.argmax(c, axis=1)

Estimator模型预测

import tensorflow as tf
# 加载模型 & 签名
imported = tf.saved_model.load(export_dir)
f = imported.signatures["predict"]

# 转换为tensor并预测
out_df = pd.DataFrame()
def predict(dataframe):
    examples = []
    for row in dataframe.itertuples():
        feature_map = {}
        # 特征处理 将特征放入dict中
        example = tf.train.Example(
            features=tf.train.Features(
                feature = feature_map
            )
        )
        examples.append(example.SerializeToString())
            
    ex = tf.constant(examples)
    result = f(examples=ex)
    out_df['high_rank_score'] = np.max(result["probabilities"].numpy(), axis=1)
    out_df['tag'] = np.argmax(result["probabilities"].numpy(), axis=1)
    return out_df