Internal cleanup.

PiperOrigin-RevId: 284259997

official/benchmark/retinanet_benchmark.py

@@ -225,6 +225,7 @@ class RetinanetBenchmarkReal(RetinanetAccuracy):
   @flagsaver.flagsaver
   def benchmark_8_gpu_coco(self):
     """Run RetinaNet model accuracy test with 8 GPUs."""
+    self.num_gpus = 8
     self._setup()
     params = copy.deepcopy(self.params_override)
     params['train']['total_steps'] = 1875  # One epoch.
@@ -236,6 +237,7 @@ class RetinanetBenchmarkReal(RetinanetAccuracy):
     # Related bug: b/135933080
     params['train']['iterations_per_loop'] = 1
     params['eval']['eval_samples'] = 8
+    FLAGS.num_gpus = self.num_gpus
     FLAGS.params_override = json.dumps(params)
     FLAGS.model_dir = self._get_model_dir('real_benchmark_8_gpu_coco')
     # Use negative value to avoid saving checkpoints.
@@ -255,6 +257,7 @@ class RetinanetBenchmarkReal(RetinanetAccuracy):
     params['train']['total_steps'] = 200
     params['train']['iterations_per_loop'] = 1
     params['eval']['eval_samples'] = 8
+    FLAGS.num_gpus = self.num_gpus
     FLAGS.params_override = json.dumps(params)
     FLAGS.model_dir = self._get_model_dir('real_benchmark_1_gpu_coco')
     FLAGS.strategy_type = 'one_device_gpu'
@@ -275,6 +278,7 @@ class RetinanetBenchmarkReal(RetinanetAccuracy):
     params['train']['total_steps'] = 200
     params['train']['iterations_per_loop'] = 1
     params['eval']['eval_samples'] = 8
+    FLAGS.num_gpus = self.num_gpus
     FLAGS.params_override = json.dumps(params)
     FLAGS.model_dir = self._get_model_dir('real_benchmark_1_gpu_coco')
     FLAGS.strategy_type = 'one_device_gpu'