模型导出
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 模型如何对外提供服务