Internal change

PiperOrigin-RevId: 306699912

Author: Allen Wang

official/vision/image_classification/classifier_trainer.py

@@ -101,8 +101,8 @@ def get_image_size_from_model(
 def _get_dataset_builders(params: base_configs.ExperimentConfig,
                           strategy: tf.distribute.Strategy,
                           one_hot: bool
-                         ) -> Tuple[Any, Any, Any]:
-  """Create and return train, validation, and test dataset builders."""
+                         ) -> Tuple[Any, Any]:
+  """Create and return train and validation dataset builders."""
   if one_hot:
     logging.warning('label_smoothing > 0, so datasets will be one hot encoded.')
   else:

official/vision/image_classification/dataset_factory.py

@@ -116,7 +116,7 @@ class DatasetConfig(base_config.Config):
   num_channels: Union[int, str] = 'infer'
   num_examples: Union[int, str] = 'infer'
   batch_size: int = 128
-  use_per_replica_batch_size: bool = False
+  use_per_replica_batch_size: bool = True
   num_devices: int = 1
   dtype: str = 'float32'
   one_hot: bool = True
@@ -185,14 +185,14 @@ def is_training(self) -> bool:
   def batch_size(self) -> int:
     """The batch size, multiplied by the number of replicas (if configured)."""
     if self.config.use_per_replica_batch_size:
-      return self.global_batch_size
+      return self.config.batch_size * self.config.num_devices
     else:
       return self.config.batch_size
 
   @property
   def global_batch_size(self):
     """The global batch size across all replicas."""
-    return self.config.batch_size * self.config.num_devices
+    return self.batch_size
 
   @property
   def num_steps(self) -> int: