Internal change

PiperOrigin-RevId: 280532315

official/vision/image_classification/README.md

@@ -139,7 +139,7 @@ python resnet_ctl_imagenet_main.py \
   --use_synthetic_data=false \
   --dtype=fp32 \
   --enable_eager=true \
-  --enable_tensorboard=false \
+  --enable_tensorboard=true \
   --distribution_strategy=tpu \
   --log_steps=50 \
   --single_l2_loss_op=true \
@@ -159,7 +159,7 @@ python resnet_ctl_imagenet_main.py \
   --use_synthetic_data=false \
   --dtype=fp32 \
   --enable_eager=true \
-  --enable_tensorboard=false \
+  --enable_tensorboard=true \
   --distribution_strategy=tpu \
   --log_steps=50 \
   --single_l2_loss_op=true \

official/vision/image_classification/resnet_ctl_imagenet_main.py

@@ -331,6 +331,11 @@ def run(flags_obj):
       train_single_step = tf.function(train_single_step)
       test_step = tf.function(test_step)
 
+    if flags_obj.enable_tensorboard:
+      summary_writer = tf.summary.create_file_writer(flags_obj.model_dir)
+    else:
+      summary_writer = None
+
     train_iter = iter(train_ds)
     time_callback.on_train_begin()
     for epoch in range(train_epochs):
@@ -371,7 +376,19 @@ def run(flags_obj):
                      test_accuracy.result().numpy(),
                      epoch + 1)
 
+      if summary_writer:
+        current_steps = steps_in_current_epoch + (epoch * per_epoch_steps)
+        with summary_writer.as_default():
+          tf.summary.scalar('train_loss', train_loss.result(), current_steps)
+          tf.summary.scalar(
+              'train_accuracy', training_accuracy.result(), current_steps)
+          tf.summary.scalar('eval_loss', test_loss.result(), current_steps)
+          tf.summary.scalar(
+              'eval_accuracy', test_accuracy.result(), current_steps)
+
     time_callback.on_train_end()
+    if summary_writer:
+      summary_writer.close()
 
     eval_result = None
     train_result = None