Changes PruningSummaries to reuse the same filewriter rather than create a new one each time, which puts tensorboard under pressure.

PiperOrigin-RevId: 318983774

Author: tensorflower-gardener

tensorflow_model_optimization/python/core/sparsity/keras/pruning_callbacks.py

@@ -103,22 +103,21 @@ def __init__(self, log_dir, update_freq='epoch', **kwargs):
 
     super(PruningSummaries, self).__init__(
         log_dir=log_dir, update_freq=update_freq, **kwargs)
+    if not compat.is_v1_apis():  # TF 2.X
+      log_dir = self.log_dir + '/metrics'
+      self._file_writer = tf.summary.create_file_writer(log_dir)
 
   def _log_pruning_metrics(self, logs, prefix, step):
     if compat.is_v1_apis():
       # Safely depend on TF 1.X private API given
       # no more 1.X releases.
       self._write_custom_summaries(step, logs)
-    else:  # TF 2.X
-      log_dir = self.log_dir + '/metrics'
-
-      file_writer = tf.summary.create_file_writer(log_dir)
-      file_writer.set_as_default()
-
-      for name, value in logs.items():
-        tf.summary.scalar(name, value, step=step)
+    else:
+      with self._file_writer.as_default():
+        for name, value in logs.items():
+          tf.summary.scalar(name, value, step=step)
 
-      file_writer.flush()
+        self._file_writer.flush()
 
   def on_epoch_begin(self, epoch, logs=None):
     if logs is not None: