allow imagenet training profiling to be disabled in config

Set option --config.profile=False on launch to disable. PiperOrigin-RevId: 647843578
google · Jun 29, 2024 · e662fbc · e662fbc
1 parent cb646c8
commit e662fbc
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/examples/imagenet/configs/default.py b/examples/imagenet/configs/default.py
@@ -56,6 +56,10 @@ def get_config():
   # num_epochs using the entire dataset. Similarly for steps_per_eval.
   config.num_train_steps = -1
   config.steps_per_eval = -1
+
+  # whether to profile the training loop
+  config.profile = True
+
   return config
 
 

diff --git a/examples/imagenet/train.py b/examples/imagenet/train.py
@@ -371,7 +371,7 @@ def train_and_evaluate(
 
   train_metrics = []
   hooks = []
-  if jax.process_index() == 0:
+  if jax.process_index() == 0 and config.profile:
     hooks += [periodic_actions.Profile(num_profile_steps=5, logdir=workdir)]
   train_metrics_last_t = time.time()
   logging.info('Initial compilation, this might take some minutes...')