From 0a2e406480d255b77fe74a2cf660b2ecccf0d859 Mon Sep 17 00:00:00 2001 From: Chanran Kim Date: Wed, 24 Nov 2021 12:51:08 +0900 Subject: [PATCH] To make the tensorboard writer compatible to tf2 To make the tensorboard writer compatible to tf2 in visualizer.py 1. version check for tf2 or not 2. FileWriter to create_file_writer and following changes --- util/visualizer.py | 31 ++++++++++++++++++++++--------- 1 file changed, 22 insertions(+), 9 deletions(-) diff --git a/util/visualizer.py b/util/visualizer.py index 584ac45d..fb836ce5 100755 --- a/util/visualizer.py +++ b/util/visualizer.py @@ -19,9 +19,13 @@ def __init__(self, opt): self.name = opt.name if self.tf_log: import tensorflow as tf + from packaging import version self.tf = tf self.log_dir = os.path.join(opt.checkpoints_dir, opt.name, 'logs') - self.writer = tf.summary.FileWriter(self.log_dir) + if version.parse(tf.__version__) >= version.parse("2"): + self.writer = tf.summary.create_file_writer(self.log_dir) + else: + self.writer = tf.summary.FileWriter(self.log_dir) if self.use_html: self.web_dir = os.path.join(opt.checkpoints_dir, opt.name, 'web') @@ -44,14 +48,19 @@ def display_current_results(self, visuals, epoch, step): except: s = BytesIO() scipy.misc.toimage(image_numpy).save(s, format="jpeg") - # Create an Image object - img_sum = self.tf.Summary.Image(encoded_image_string=s.getvalue(), height=image_numpy.shape[0], width=image_numpy.shape[1]) - # Create a Summary value - img_summaries.append(self.tf.Summary.Value(tag=label, image=img_sum)) + if self.version.parse(self.tf.__version__) >= self.version.parse("2"): + with self.writer.as_default(): + self.tf.summary.image(label, np.expand_dims(image_numpy, axis=0), step) + else: + # Create an Image object + img_sum = self.tf.Summary.Image(encoded_image_string=s.getvalue(), height=image_numpy.shape[0], width=image_numpy.shape[1]) + # Create a Summary value + img_summaries.append(self.tf.Summary.Value(tag=label, image=img_sum)) # Create and write Summary - summary = self.tf.Summary(value=img_summaries) - self.writer.add_summary(summary, step) + if self.version.parse(self.tf.__version__) < self.version.parse("2"): + summary = self.tf.Summary(value=img_summaries) + self.writer.add_summary(summary, step) if self.use_html: # save images to a html file for label, image_numpy in visuals.items(): @@ -95,8 +104,12 @@ def display_current_results(self, visuals, epoch, step): def plot_current_errors(self, errors, step): if self.tf_log: for tag, value in errors.items(): - summary = self.tf.Summary(value=[self.tf.Summary.Value(tag=tag, simple_value=value)]) - self.writer.add_summary(summary, step) + if self.version.parse(self.tf.__version__) >= self.version.parse("2"): + with self.writer.as_default(): + self.tf.summary.scalar(tag, value, step) + else: + summary = self.tf.Summary(value=[self.tf.Summary.Value(tag=tag, simple_value=value)]) + self.writer.add_summary(summary, step) # errors: same format as |errors| of plotCurrentErrors def print_current_errors(self, epoch, i, errors, t):