Estimator.train() and .predict() are too slow for small data sets
Convert to a tf.keras.Model instead of an Estimator, and use tf.keras.Model.fit() instead of Estimator.train(). fit() doesn't have the fixed delay that train() does. The Keras predict() doesn't either.
I'm trying to implement a DQN which makes many calls to Estimator.train() followed by Estimator.predict() on the same model with a small number of examples each. But each call takes a minimum of a few hundred milliseconds to over a second which is independent of the number of examples for small numbers like 1-20.
I think these delays are caused by rebuilding the graph and saving checkpoints on each call. Is there are way to keep the same graph and parameters in memory for fast train-predict iterations or otherwise speed it up?