Internal change

PiperOrigin-RevId: 311626933

official/vision/image_classification/resnet/common.py

@@ -117,7 +117,8 @@ def get_callbacks(
 
   if FLAGS.enable_tensorboard:
     tensorboard_callback = tf.keras.callbacks.TensorBoard(
-        log_dir=FLAGS.model_dir)
+        log_dir=FLAGS.model_dir,
+        profile_batch=FLAGS.profile_steps)
     callbacks.append(tensorboard_callback)
 
   is_pruning_enabled = pruning_method is not None
@@ -228,6 +229,14 @@ def define_keras_flags(
   flags.DEFINE_boolean(
       name='enable_tensorboard', default=False,
       help='Whether to enable Tensorboard callback.')
+  flags.DEFINE_string(
+      name='profile_steps', default=None,
+      help='Save profiling data to model dir at given range of global steps. The '
+      'value must be a comma separated pair of positive integers, specifying '
+      'the first and last step to profile. For example, "--profile_steps=2,4" '
+      'triggers the profiler to process 3 steps, starting from the 2nd step. '
+      'Note that profiler has a non-trivial performance overhead, and the '
+      'output file can be gigantic if profiling many steps.')
   flags.DEFINE_integer(
       name='train_steps', default=None,
       help='The number of steps to run for training. If it is larger than '