Internal change

PiperOrigin-RevId: 438897452

official/core/config_definitions.py

@@ -76,6 +76,10 @@ class DataConfig(base_config.Config):
       features. The main use case is to skip the image/video decoding for better
       performance.
     seed: An optional seed to use for deterministic shuffling/preprocessing.
+    prefetch_buffer_size: An int specifying the buffer size of prefetch
+      datasets. If None, the buffer size is autotuned. Specifying this is useful
+      in case autotuning uses up too much memory by making the buffer size too
+      high.
   """
   input_path: Union[Sequence[str], str, base_config.Config] = ""
   tfds_name: str = ""
@@ -96,6 +100,7 @@ class DataConfig(base_config.Config):
   tfds_as_supervised: bool = False
   tfds_skip_decoding_feature: str = ""
   seed: Optional[int] = None
+  prefetch_buffer_size: Optional[int] = None
 
 
 @dataclasses.dataclass

official/core/input_reader.py

@@ -292,6 +292,8 @@ class InputReader:
     self._transform_and_batch_fn = transform_and_batch_fn
     self._postprocess_fn = postprocess_fn
     self._seed = params.seed
+    self._prefetch_buffer_size = (params.prefetch_buffer_size or
+                                  tf.data.experimental.AUTOTUNE)
 
     # When tf.data service is enabled, each data service worker should get
     # different random seeds. Thus, we set `seed` to None.
@@ -505,4 +507,4 @@ class InputReader:
       options = tf.data.Options()
       options.experimental_deterministic = self._deterministic
       dataset = dataset.with_options(options)
-    return dataset.prefetch(tf.data.experimental.AUTOTUNE)
+    return dataset.prefetch(self._prefetch_buffer_size)

official/vision/configs/experiments/image_classification/imagenet_resnet50_gpu.yaml

@@ -19,12 +19,16 @@ task:
     is_training: true
     global_batch_size: 2048
     dtype: 'float16'
+    # Autotuning the prefetch buffer size causes OOMs, so set it to a reasonable
+    # static value: 32. See b/218880025.
+    prefetch_buffer_size: 32
   validation_data:
     input_path: 'imagenet-2012-tfrecord/valid*'
     is_training: false
     global_batch_size: 2048
     dtype: 'float16'
     drop_remainder: false
+    prefetch_buffer_size: 32
 trainer:
   train_steps: 56160
   validation_steps: 25