Internal change

PiperOrigin-RevId: 304805298

official/benchmark/keras_imagenet_benchmark.py

@@ -180,7 +180,8 @@ class Resnet50KerasAccuracy(keras_benchmark.KerasBenchmark):
 
   def benchmark_8_gpu(self):
     """Tests Keras model with eager, dist_strat and 8 GPUs."""
-    self._run_benchmark(
+    self._setup()
+    self._run_and_report_benchmark(
         experiment_name='benchmark_8_gpu',
         num_gpus=8,
         per_replica_batch_size=128,
@@ -191,7 +192,8 @@ class Resnet50KerasAccuracy(keras_benchmark.KerasBenchmark):
 
   def benchmark_8_gpu_fp16(self):
     """Tests Keras model with eager, dist_strat, 8 GPUs, and fp16."""
-    self._run_benchmark(
+    self._setup()
+    self._run_and_report_benchmark(
         experiment_name='benchmark_8_gpu_fp16',
         num_gpus=8,
         per_replica_batch_size=256,
@@ -202,7 +204,8 @@ class Resnet50KerasAccuracy(keras_benchmark.KerasBenchmark):
 
   def benchmark_xla_8_gpu_fp16(self):
     """Tests Keras model with XLA, eager, dist_strat, 8 GPUs and fp16."""
-    self._run_benchmark(
+    self._setup()
+    self._run_and_report_benchmark(
         experiment_name='benchmark_xla_8_gpu_fp16',
         num_gpus=8,
         per_replica_batch_size=256,
@@ -214,7 +217,8 @@ class Resnet50KerasAccuracy(keras_benchmark.KerasBenchmark):
 
   def benchmark_xla_8_gpu_fp16_dynamic(self):
     """Tests Keras model with XLA, eager, dist_strat, 8 GPUs, dynamic fp16."""
-    self._run_benchmark(
+    self._setup()
+    self._run_and_report_benchmark(
         experiment_name='benchmark_xla_8_gpu_fp16_dynamic',
         top_1_min=0.736,
         num_gpus=8,

official/vision/image_classification/classifier_trainer_test.py

@@ -340,8 +340,7 @@ class UtilTests(parameterized.TestCase, tf.test.TestCase):
             dataset_num_private_threads=1,
         ),
         train_dataset=dataset_factory.DatasetConfig(dtype=dtype),
-        model=base_configs.ModelConfig(
-            loss=base_configs.LossConfig(loss_scale='dynamic')),
+        model=base_configs.ModelConfig(),
     )
 
     class EmptyClass:

official/vision/image_classification/configs/base_configs.py

@@ -117,13 +117,11 @@ class LossConfig(base_config.Config):
 
   Attributes:
     name: The name of the loss. Defaults to None.
-    loss_scale: The type of loss scale
     label_smoothing: Whether or not to apply label smoothing to the loss. This
       only applies to 'categorical_cross_entropy'.
 
   """
   name: str = None
-  loss_scale: str = None
   label_smoothing: float = None