Add caching to input pipeline (#34)

Caching the results pre-shuffling prevents TF from recomputing the pre-shuffle batches from scratch repeatedly.
A speedup of about >5x is expected for the input pipeline, and similarly for the training cycle times as well.

Author: Northbadge

compiler_opt/rl/data_reader.py

@@ -165,15 +165,16 @@ def create_sequence_example_dataset_fn(
   def _sequence_example_dataset_fn(sequence_examples):
     # Data collector returns empty strings for corner cases, filter them out
     # here.
-    dataset = tf.data.Dataset.from_tensor_slices(sequence_examples).filter(
-        lambda string: tf.strings.length(string) > 0).map(parser_fn).filter(
-            lambda traj: tf.size(traj.reward) > 2)
-    dataset = (
-        dataset.unbatch().batch(
-            train_sequence_length,
-            drop_remainder=True).shuffle(trajectory_shuffle_buffer_size).batch(
-                batch_size,
-                drop_remainder=True).prefetch(tf.data.experimental.AUTOTUNE))
+    dataset = (tf.data.Dataset.from_tensor_slices(sequence_examples)
+              .filter(lambda string: tf.strings.length(string) > 0)
+              .map(parser_fn)
+              .filter(lambda traj: tf.size(traj.reward) > 2)
+              .unbatch()
+              .batch(train_sequence_length, drop_remainder=True)
+              .cache()
+              .shuffle(trajectory_shuffle_buffer_size)
+              .batch(batch_size, drop_remainder=True)
+              )
     return dataset
 
   return _sequence_example_dataset_fn

compiler_opt/rl/train_locally.py

@@ -126,7 +126,7 @@ def train_eval(agent_name=constant.AgentName.PPO,
       train_sequence_length=train_sequence_length)
 
   def sequence_example_iterator_fn(seq_ex: List[str]):
-    return iter(dataset_fn(seq_ex).repeat())
+    return iter(dataset_fn(seq_ex).repeat().prefetch(tf.data.AUTOTUNE))
 
   reward_stat_map = collections.defaultdict(lambda: None)
   reward_stat_map_path = os.path.join(root_dir, 'reward_stat_map')