在等待实时预测输入时将 tensorflow 估计器保存在内存中

Question

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

在代码的开头我实例化了估计器：

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] Creating TensorFlow device (/device:GPU:0) -> (device: 0, name: GeForce GTX 1060 6GB, pci bus id: 0000:04:00.0, compute capability: 6.1)

INFO:tensorflow:Restoring parameters from /home/fgervais/tf/model_dir_3/model.ckpt-103712

似乎整个 GPU 检测过程和模型加载在每个预测上都重新完成。

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

Answer 1

解决这个问题的方法是使用 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)})

在等待实时预测输入时将 tensorflow 估计器保存在内存中

Keep the tensorflow estimator in memory while waiting for live prediction inputs

tensorflow

tensorflow-datasets