allow imagenet training profiling to be disabled in config

Set option --config.profile=False on launch to disable.

PiperOrigin-RevId: 647843578

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
 
 

examples/imagenet/train.py

@@ -371,8 +371,8 @@ def train_and_evaluate(
 
   train_metrics = []
   hooks = []
-  if jax.process_index() == 0:
-    hooks += [periodic_actions.Profile(num_profile_steps=5, logdir=workdir)]
+  if jax.process_index() == 0 and config.profile:
+    hooks += [periodic_actions.Profile(num_profile_steps=3, logdir=workdir)]
   train_metrics_last_t = time.time()
   logging.info('Initial compilation, this might take some minutes...')
   for step, batch in zip(range(step_offset, num_steps), train_iter):