并行数据提取

eda/esmm/Model_pipline/train.py

@@ -81,8 +81,15 @@ def input_fn(filenames, batch_size=32, num_epochs=1, perform_shuffle=False):
     # dataset = dataset.repeat(num_epochs)
     # dataset = dataset.batch(batch_size) # Batch size to use
 
-
-    dataset = tf.data.TFRecordDataset(filenames).apply(tf.data.experimental.map_and_batch(map_func=_parse_fn, batch_size=batch_size, num_parallel_calls=8))
+    files = tf.data.Dataset.list_files(filenames)
+    dataset = files.apply(
+        tf.data.experimental.parallel_interleave(
+            lambda  file: tf.data.TFRecordDataset(file),
+            cycle_length=8
+        )
+    )
+
+    dataset = dataset.apply(tf.data.experimental.map_and_batch(map_func=_parse_fn, batch_size=batch_size, num_parallel_calls=8))
     dataset = dataset.prefetch(10000)