measure first batch end time to estimate start_time = first_batch_end - avg_step_time.

PiperOrigin-RevId: 331795709

official/utils/misc/keras_utils.py

@@ -30,9 +30,13 @@ from tensorflow.python.eager import monitoring
 global_batch_size_gauge = monitoring.IntGauge(
     '/tensorflow/training/global_batch_size', 'TF training global batch size')
 
-first_batch_start_time = monitoring.IntGauge(
-    '/tensorflow/training/first_batch_start',
-    'TF training start time (unix epoch time in us.')
+first_batch_time_gauge = monitoring.IntGauge(
+    '/tensorflow/training/first_batch',
+    'TF training start/end time for first batch (unix epoch time in us.',
+    'type')
+
+first_batch_start_time = first_batch_time_gauge.get_cell('start')
+first_batch_end_time = first_batch_time_gauge.get_cell('end')
 
 
 class BatchTimestamp(object):
@@ -121,8 +125,8 @@ class TimeHistory(tf.keras.callbacks.Callback):
   def on_batch_begin(self, batch, logs=None):
     if not self.start_time:
       self.start_time = time.time()
-      if not first_batch_start_time.get_cell().value():
-        first_batch_start_time.get_cell().set(int(self.start_time * 1000000))
+      if not first_batch_start_time.value():
+        first_batch_start_time.set(int(self.start_time * 1000000))
 
     # Record the timestamp of the first global step
     if not self.timestamp_log:
@@ -131,6 +135,8 @@ class TimeHistory(tf.keras.callbacks.Callback):
 
   def on_batch_end(self, batch, logs=None):
     """Records elapse time of the batch and calculates examples per second."""
+    if not first_batch_end_time.value():
+      first_batch_end_time.set(int(time.time() * 1000000))
     self.steps_in_epoch = batch + 1
     steps_since_last_log = self.global_steps - self.last_log_step
     if steps_since_last_log >= self.log_steps: