在等待实时预测输入时将张量流估计器保留在内存中

时间:2018-04-22 13:30:22

标签: tensorflow tensorflow-datasets

我有一个训练有素的估算器,我用它做新的输入数据时进行实时预测。

在代码的开头,我实例化了估算器:

estimator = tf.estimator.Estimator(
    model_fn=model_fn,
    model_dir="{}/model_dir_{}".format(script_dir, 3))

然后在循环中,每次我获得足够的新数据进行预测时我都会这样做:

predict_input_fn = tf.estimator.inputs.numpy_input_fn(
    x={"x": np.array([sample.normalized.input_data])},
    num_epochs=1,
    shuffle=False)
predictions = estimator.predict(
    input_fn=predict_input_fn,
)

每次我这样做,我都会在控制台中收到这些tensorflow消息:

  

2018-04-21 16:01:08.401319:I tensorflow / core / common_runtime / gpu / gpu_device.cc:1195]创建TensorFlow设备(/ device:GPU:0) - > (设备:0,名称:GeForce GTX 1060 6GB,pci总线ID:0000:04:00.0,计算能力:6.1)

     

INFO:tensorflow:从/home/fgervais/tf/model_dir_3/model.ckpt-103712恢复参数

似乎整个GPU检测过程和模型加载都会在每次预测时再次完成。

有没有办法让模型在实时输入之间加载到内存中,以便获得更好的预测率?

1 个答案:

答案 0 :(得分:3)

解决方法是使用predictor

在问题的具体情况下,可以这样做:

def serving_input_fn():
    x = tf.placeholder(dtype=tf.float32, shape=[3500], name='x')
    inputs = {'x': x }

    return tf.estimator.export.ServingInputReceiver(inputs, inputs)

estimator = tf.estimator.Estimator(
    model_fn=model_fn,
    model_dir="{}/model_dir_{}/model.ckpt-103712".format(script_dir, 3))

estimator_predictor = tf.contrib.predictor.from_estimator(
                            estimator, serving_input_fn)

p = estimator_predictor(
        {"x": np.array(sample.normalized.input_data)})