[tensorflow]TensorFlow模型如何对外提供服务

save函数实现(关键代码)

def raw_infer_fn(self):
    c_size = self.params["field_size"]
    features = {
        'feat_ids': tf.placeholder(tf.int64, [None, c_size], name='feat_ids'),
        'feat_vals': tf.placeholder(tf.float32, [None, c_size], name='feat_vals')
    }
    s = tf.estimator.export.build_raw_serving_input_receiver_fn(features)
    return s
def save(self, model_dir, input_type):
        base_dir = "./multi_model"
        full_dir = os.path.join(base_dir, model_dir)
        if input_type == "proto":
            self.model.export_saved_model(full_dir, self.proto_infer_fn())
        else:
            self.model.export_saved_model(full_dir, self.raw_infer_fn())

调用方式

model.save("push_rank", "raw")

样本数据 df_tmp

feat_ids, feat_vals = [], []
for i, cur in enumerate(df_tmp.iloc[0].tolist()):
    if i < 2:
        continue
    ids, vals = cur.split(':')
    feat_ids.append(int(ids))
    feat_vals.append(vals)
input_dict = {"feat_ids":[feat_ids], "feat_vals":[feat_vals]}
print(input_dict)

from tensorflow import contrib
# 从导出目录中加载模型，并生成预测函数。
predict_fn = tf.contrib.predictor.from_saved_model('./multi_model/push_rank/2')
# 开始预测
#input_dict = {"feat_ids":[[1]*816], "feat_vals":[[1]*816]}
predictions = predict_fn(input_dict)
print(predictions)

out

{'prob2': array([[0.2861896]], dtype=float32), 'prob1': array([[0.2861896]], dtype=float32)}

References
TensorFlow 模型如何对外提供服务