performance - Tensorflow Estimator の予測が遅い

Question

tf.estimator.LinearClassifier をトレーニングしました。モデルのトレーニングと評価にはデータ サイズに対して妥当な時間 (~60 秒) がかかりますが、予測には桁違いに長い時間がかかります (~1 時間)。

予測コードは次のとおりです。

predictionResult = estimator.predict(input_fn=lambda: my_input_fn2(predictionValidationFile, False, 1))
predictionList = [prediction for prediction in  predictionResult]

と：

def my_input_fn2(file_path, perform_shuffle=False, repeat_count=1):
def _parse_function(example_proto):      
  keys_to_features = {"xslm": tf.FixedLenFeature([10000], tf.float32),
                      "xrnn": tf.FixedLenFeature([10000], tf.float32),
                      "target": tf.FixedLenFeature([10000], tf.float32)}
  parsed_features = tf.parse_single_example(example_proto, keys_to_features)      
  myfeatures = {'xrnn':parsed_features['xrnn'], 'xslm':parsed_features['xslm']}
  return myfeatures, parsed_features['target'] 

dataset = (tf.data.TFRecordDataset(file_path)                
           .map(_parse_function))     
dataset = dataset.repeat(repeat_count) 
dataset = dataset.batch(1)  
iterator = dataset.make_one_shot_iterator()
batch_feature,  batch_labels = iterator.get_next()    
xs= tf.reshape(batch_feature['xslm'],[-1,1])
xr= tf.reshape(batch_feature['xrnn'],[-1,1])
x = {'xrnn':xr, 'xslm':xs}
y = tf.reshape(batch_labels, [-1,1])
return x, y

2 行目は、10,000 サンプル (1 バッチに相当) 実行した場合、実行に 0.8 秒かかります。50 000 000 サンプルの場合、予測には 1 時間以上かかります。

この段階での私の推測では、このパフォーマンスの低下は、エスティメーターの predict() 関数が実際の予測結果を返すのではなく、python ジェネレーターを返しているという事実によって単純に引き起こされているということです。バッチごとに、ジェネレーターは最終的に関数を 10,000 回呼び出して、10,000 個の予測結果を取得します。これは非効率に思えます。

スピードアップするオプションはありますか？