[Benchmarking-Py] Replacing tf.data.experimental.AUTOTUNE with tf.data.AUTOTUNE

Author: DEKHTIARJonathan

tftrt/benchmarking-python/dataloading_utils.py

@@ -12,7 +12,7 @@ def SyntheticDataset(dataset, device):
     dataset = dataset.take(count=1)  # loop over 1 batch
     dataset = dataset.cache()
     dataset = dataset.repeat()
-    dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+    dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
     dataset = ensure_dataset_on_gpu(dataset, device)
     return dataset
 
@@ -46,7 +46,7 @@ def ensure_dataset_on_gpu(dataset, device):
         print(f"[INFO] Adding prefetch to device `{device}` to the dataset.")
         dataset = dataset.apply(
             tf.data.experimental.prefetch_to_device(
-                device=device, buffer_size=tf.data.experimental.AUTOTUNE
+                device=device, buffer_size=tf.data.AUTOTUNE
             )
         )
         return dataset

tftrt/benchmarking-python/huggingface/bert/transformers.py

@@ -128,7 +128,7 @@ def get_dataset_batches(self):
         dataset = dataset.take(count=1)  # loop over 1 batch
         dataset = dataset.cache()
         dataset = dataset.repeat()
-        dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
+        dataset = dataset.prefetch(tf.data.AUTOTUNE)
 
         return dataset, None
 

tftrt/benchmarking-python/huggingface/t5/infer.py

@@ -169,7 +169,7 @@ def get_dataset_batches(self):
             dataset = dataset.cache()
             dataset = dataset.repeat()
 
-        dataset = dataset.prefetch(tf.data.experimental.AUTOTUNE)
+        dataset = dataset.prefetch(tf.data.AUTOTUNE)
         return dataset, None
 
     def preprocess_model_inputs(self, data_batch):

tftrt/benchmarking-python/image_classification/image_classification.py

@@ -163,7 +163,7 @@ def preprocess_sample_fn(record):
 
         dataset = dataset.interleave(
             tf.data.TFRecordDataset,
-            cycle_length=tf.data.experimental.AUTOTUNE,
+            cycle_length=tf.data.AUTOTUNE,
             block_length=max(self._args.batch_size, 32)
         )
 
@@ -175,12 +175,12 @@ def preprocess_sample_fn(record):
 
         dataset = dataset.map(
             map_func=preprocess_fn,
-            num_parallel_calls=tf.data.experimental.AUTOTUNE,
+            num_parallel_calls=tf.data.AUTOTUNE,
         )
 
         dataset = dataset.batch(self._args.batch_size, drop_remainder=False)
 
-        dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+        dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
 
         return dataset, None
 

tftrt/benchmarking-python/nvidia_examples/bert_tf2/infer.py

@@ -200,7 +200,7 @@ def append_feature(feature, is_padding):
         )
 
         dataset = dataset.batch(self._args.batch_size, drop_remainder=False)
-        dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+        dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
 
         bypass_data_to_eval = {
             "eval_features": eval_features,

tftrt/benchmarking-python/nvidia_examples/efficientnet_base/infer.py

@@ -196,17 +196,17 @@ def parse_record(record, image_size):
             tf.data.TFRecordDataset,
             cycle_length=10,
             block_length=1,
-            num_parallel_calls=tf.data.experimental.AUTOTUNE
+            num_parallel_calls=tf.data.AUTOTUNE
         )
 
         parse_record_fn = lambda record: parse_record(
             record=record, image_size=self._args.input_size
         )
         dataset = dataset.map(
-            parse_record_fn, num_parallel_calls=tf.data.experimental.AUTOTUNE
+            parse_record_fn, num_parallel_calls=tf.data.AUTOTUNE
         )
         dataset = dataset.batch(self._args.batch_size, drop_remainder=False)
-        dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+        dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
 
         return dataset, None
 

tftrt/benchmarking-python/nvidia_examples/electra_tf2/infer.py

@@ -220,9 +220,7 @@ def get_dataset_from_features(features, batch_size):
             ))
 
             dataset = dataset.batch(batch_size, drop_remainder=False)
-            dataset = dataset.prefetch(
-                buffer_size=tf.data.experimental.AUTOTUNE
-            )
+            dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
 
             return dataset
 

tftrt/benchmarking-python/nvidia_examples/mrcnn_tf2/infer.py

@@ -167,14 +167,14 @@ def get_dataset_batches(self):
 
         dataset = dataset.map(
             lambda x: dataset_parser(value=x, args=self._args),
-            num_parallel_calls=tf.data.experimental.AUTOTUNE
+            num_parallel_calls=tf.data.AUTOTUNE
         )
 
         dataset = dataset.batch(
             batch_size=args.batch_size, drop_remainder=False
         )
 
-        dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+        dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
 
         return dataset, None
 

tftrt/benchmarking-python/nvidia_examples/sim_tf2/infer.py

@@ -108,16 +108,16 @@ def _remap_values(sample, out_type):
 
         dataset = dataset.map(
             map_func=partial(tf.io.parse_example, features=feature_spec),
-            num_parallel_calls=tf.data.experimental.AUTOTUNE
+            num_parallel_calls=tf.data.AUTOTUNE
         )
 
         out_type = tf.float16 if self._args.amp else tf.float32
         dataset = dataset.map(
             map_func=partial(_remap_values, out_type=out_type),
-            num_parallel_calls=tf.data.experimental.AUTOTUNE
+            num_parallel_calls=tf.data.AUTOTUNE
         )
 
-        dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+        dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
 
         return dataset, None
 

tftrt/benchmarking-python/nvidia_examples/unet_medical_tf2/infer.py

@@ -101,7 +101,7 @@ def get_dataset_batches(self):
 
         dataset = dataset.batch(args.batch_size, drop_remainder=False)
 
-        dataset = dataset.prefetch(buffer_size=tf.data.experimental.AUTOTUNE)
+        dataset = dataset.prefetch(buffer_size=tf.data.AUTOTUNE)
 
         return dataset, None