Skip to main content
 首页 » 编程设计

python之如何通过 export_savedmodel 保存一个 tensorflow estimator,然后在本地加载和使用它

2024年10月01日37sky-heaven

我训练了一个 tf 线性回归估计器,如下所示:

sample_size = train_x.shape[0] 
feature_size = train_x.shape[1] 
 
feature_columns = [tf.feature_column.numeric_column("x", shape=[feature_size])] 
 
lr_estimator = tf.estimator.LinearRegressor(feature_columns=feature_columns ) 
 
train_x_mat = train_x.as_matrix() 
test_x_mat = test_x.as_matrix() 
 
 
# Define the training inputs 
train_input_fn = tf.estimator.inputs.numpy_input_fn( 
    x={"x": train_x_mat}, 
    y=np.array(train_y_mat), 
    num_epochs=None, 
    shuffle=True) 
 
# Train model. 
lr_estimator.train(input_fn=train_input_fn, steps=2000) 

其中 train_x 和 train_y 是 pandas 数据框。 lr_estimator 确实有效,我可以成功调用 .predict。

我如何将它保存到一个文件中,然后加载回来以供稍后预测?我只是想构建一个小的 python 程序。预测程序将在同一个桌面上运行。我还不需要复杂的服务器服务。

请您参考如下方法:

def serving_input_receiver_fn(): 
    """ 
    input placeholder 
    """ 
    inputs = {"x": tf.placeholder(shape=[feature_size], dtype=tf.float32)} 
    return tf.estimator.export.ServingInputReceiver(inputs, inputs) 
 
# export model and weights 
export_dir = est_inception_v3.export_savedmodel(export_dir_base="/export_dir",  
    serving_input_receiver_fn=serving_input_receiver_fn) 
 
# restore from disk 
with tf.Session() as sess: 
    tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], export_dir) 
    predictor = SavedModelPredictor(export_dir) 
    print(predictor({"x": test_x_mat}))