Merge pull request #208 from NREL/bnb/tensorboard_logging

tensorboard logging

sup3r/models/abstract.py

@@ -1,7 +1,5 @@
 # -*- coding: utf-8 -*-
-"""
-Abstract class to define the required interface for Sup3r model subclasses
-"""
+"""Abstract class defining the required interface for Sup3r model subclasses"""
 import json
 import logging
 import os
@@ -23,10 +21,83 @@
 import sup3r.utilities.loss_metrics
 from sup3r.preprocessing.data_handling.exogenous_data_handling import ExoData
 from sup3r.utilities import VERSION_RECORD
+from sup3r.utilities.utilities import Timer
 
 logger = logging.getLogger(__name__)
 
 
+class TensorboardMixIn:
+    """MixIn class for tensorboard logging and profiling."""
+
+    def __init__(self):
+        self._tb_writer = None
+        self._tb_log_dir = None
+        self._write_tb_profile = False
+        self._total_batches = None
+        self._history = None
+        self.timer = Timer()
+
+    @property
+    def total_batches(self):
+        """Record of total number of batches for logging."""
+        if self._total_batches is None and self._history is None:
+            self._total_batches = 0
+        elif self._history is None and 'total_batches' in self._history:
+            self._total_batches = self._history['total_batches'].values[-1]
+        elif self._total_batches is None and self._history is not None:
+            self._total_batches = 0
+        return self._total_batches
+
+    @total_batches.setter
+    def total_batches(self, value):
+        """Set total number of batches."""
+        self._total_batches = value
+
+    def dict_to_tensorboard(self, entry):
+        """Write data to tensorboard log file. This is usually a loss_details
+        dictionary.
+
+        Parameters
+        ----------
+        entry: dict
+            Dictionary of values to write to tensorboard log file
+        """
+        if self._tb_writer is not None:
+            with self._tb_writer.as_default():
+                for name, value in entry.items():
+                    if isinstance(value, str):
+                        tf.summary.text(name, value, self.total_batches)
+                    else:
+                        tf.summary.scalar(name, value, self.total_batches)
+
+    def profile_to_tensorboard(self, name):
+        """Write profile data to tensorboard log file.
+
+        Parameters
+        ----------
+        name : str
+            Tag name to use for profile info
+        """
+        if self._tb_writer is not None and self._write_tb_profile:
+            with self._tb_writer.as_default():
+                tf.summary.trace_export(name=name, step=self.total_batches,
+                                        profiler_outdir=self._tb_log_dir)
+
+    def _init_tensorboard_writer(self, out_dir):
+        """Initialize the ``tf.summary.SummaryWriter`` to use for writing
+        tensorboard compatible log files.
+
+        Parameters
+        ----------
+        out_dir : str
+            Standard out_dir where model epochs are saved. e.g. './gan_{epoch}'
+        """
+        tb_log_pardir = os.path.abspath(os.path.join(out_dir, os.pardir))
+        self._tb_log_dir = os.path.join(tb_log_pardir, 'logs')
+        os.makedirs(self._tb_log_dir, exist_ok=True)
+        self._tb_writer = tf.summary.create_file_writer(self._tb_log_dir)
+
+
 class AbstractInterface(ABC):
     """
     Abstract class to define the required interface for Sup3r model subclasses
@@ -371,9 +442,8 @@ def hr_exo_features(self):
         # pylint: disable=E1101
         features = []
         if hasattr(self, '_gen'):
-            for layer in self._gen.layers:
-                if isinstance(layer, (Sup3rAdder, Sup3rConcat)):
-                    features.append(layer.name)
+            features = [layer.name for layer in self._gen.layers
+                        if isinstance(layer, (Sup3rAdder, Sup3rConcat))]
         return features
 
     @property
@@ -465,13 +535,14 @@ def save_params(self, out_dir):
 
 
 # pylint: disable=E1101,W0201,E0203
-class AbstractSingleModel(ABC):
+class AbstractSingleModel(ABC, TensorboardMixIn):
     """
     Abstract class to define the required training interface
     for Sup3r model subclasses
     """
 
     def __init__(self):
+        super().__init__()
         self.gpu_list = tf.config.list_physical_devices('GPU')
         self.default_device = '/cpu:0'
         self._version_record = VERSION_RECORD
@@ -743,13 +814,13 @@ def init_optimizer(optimizer, learning_rate):
         """
         if isinstance(optimizer, dict):
             class_name = optimizer['name']
-            OptimizerClass = getattr(optimizers, class_name)
-            sig = signature(OptimizerClass)
+            optimizer_class = getattr(optimizers, class_name)
+            sig = signature(optimizer_class)
             optimizer_kwargs = {
                 k: v
                 for k, v in optimizer.items() if k in sig.parameters
             }
-            optimizer = OptimizerClass.from_config(optimizer_kwargs)
+            optimizer = optimizer_class.from_config(optimizer_kwargs)
         elif optimizer is None:
             optimizer = optimizers.Adam(learning_rate=learning_rate)
 
@@ -915,10 +986,9 @@ def update_loss_details(loss_details, new_data, batch_len, prefix=None):
         prior_n_obs = loss_details['n_obs']
         new_n_obs = prior_n_obs + batch_len
 
-        for key, new_value in new_data.items():
-            key = key if prefix is None else prefix + key
-            new_value = (new_value if not isinstance(new_value, tf.Tensor) else
-                         new_value.numpy())
+        for k, v in new_data.items():
+            key = k if prefix is None else prefix + k
+            new_value = (v if not isinstance(v, tf.Tensor) else v.numpy())
 
             if key in loss_details:
                 saved_value = loss_details[key]
@@ -1061,9 +1131,7 @@ def finish_epoch(self,
         stop : bool
             Flag to early stop training.
         """
-
         self.log_loss_details(loss_details)
-
         self._history.at[epoch, 'elapsed_time'] = time.time() - t0
         for key, value in loss_details.items():
             if key != 'n_obs':
@@ -1135,20 +1203,19 @@ def run_gradient_descent(self,
         loss_details : dict
             Namespace of the breakdown of loss components
         """
-
         t0 = time.time()
         if optimizer is None:
             optimizer = self.optimizer
 
         if not multi_gpu or len(self.gpu_list) == 1:
+
             grad, loss_details = self.get_single_grad(low_res, hi_res_true,
                                                       training_weights,
                                                       **calc_loss_kwargs)
             optimizer.apply_gradients(zip(grad, training_weights))
             t1 = time.time()
             logger.debug(f'Finished single gradient descent step '
                          f'in {(t1 - t0):.3f}s')
-
         else:
             futures = []
             lr_chunks = np.array_split(low_res, len(self.gpu_list))
@@ -1178,7 +1245,6 @@ def run_gradient_descent(self,
             t1 = time.time()
             logger.debug(f'Finished {len(futures)} gradient descent steps on '
                          f'{len(self.gpu_list)} GPUs in {(t1 - t0):.3f}s')
-
         return loss_details
 
     def _reshape_norm_exo(self, hi_res, hi_res_exo, exo_name, norm_in=True):
@@ -1283,8 +1349,10 @@ def generate(self,
             low_res = self.norm_input(low_res)
 
         hi_res = self.generator.layers[0](low_res)
-        for i, layer in enumerate(self.generator.layers[1:]):
-            try:
+        layer_num = 1
+        try:
+            for i, layer in enumerate(self.generator.layers[1:]):
+                layer_num = i + 1
                 if isinstance(layer, (Sup3rAdder, Sup3rConcat)):
                     msg = (f'layer.name = {layer.name} does not match any '
                            'features in exogenous_data '
@@ -1299,11 +1367,11 @@ def generate(self,
                     hi_res = layer(hi_res, hi_res_exo)
                 else:
                     hi_res = layer(hi_res)
-            except Exception as e:
-                msg = ('Could not run layer #{} "{}" on tensor of shape {}'.
-                       format(i + 1, layer, hi_res.shape))
-                logger.error(msg)
-                raise RuntimeError(msg) from e
+        except Exception as e:
+            msg = ('Could not run layer #{} "{}" on tensor of shape {}'.
+                   format(layer_num, layer, hi_res.shape))
+            logger.error(msg)
+            raise RuntimeError(msg) from e
 
         hi_res = hi_res.numpy()
 
@@ -1341,8 +1409,10 @@ def _tf_generate(self, low_res, hi_res_exo=None):
             Synthetically generated high-resolution data
         """
         hi_res = self.generator.layers[0](low_res)
-        for i, layer in enumerate(self.generator.layers[1:]):
-            try:
+        layer_num = 1
+        try:
+            for i, layer in enumerate(self.generator.layers[1:]):
+                layer_num = i + 1
                 if isinstance(layer, (Sup3rAdder, Sup3rConcat)):
                     msg = (f'layer.name = {layer.name} does not match any '
                            f'features in exogenous_data ({list(hi_res_exo)})')
@@ -1351,11 +1421,11 @@ def _tf_generate(self, low_res, hi_res_exo=None):
                     hi_res = layer(hi_res, hr_exo)
                 else:
                     hi_res = layer(hi_res)
-            except Exception as e:
-                msg = ('Could not run layer #{} "{}" on tensor of shape {}'.
-                       format(i + 1, layer, hi_res.shape))
-                logger.error(msg)
-                raise RuntimeError(msg) from e
+        except Exception as e:
+            msg = ('Could not run layer #{} "{}" on tensor of shape {}'.
+                   format(layer_num, layer, hi_res.shape))
+            logger.error(msg)
+            raise RuntimeError(msg) from e
 
         return hi_res
 
@@ -1398,16 +1468,13 @@ def get_single_grad(self,
         loss_details : dict
             Namespace of the breakdown of loss components
         """
-        with tf.device(device_name):
-            with tf.GradientTape(watch_accessed_variables=False) as tape:
-                tape.watch(training_weights)
-
-                hi_res_exo = self.get_high_res_exo_input(hi_res_true)
-                hi_res_gen = self._tf_generate(low_res, hi_res_exo)
-                loss_out = self.calc_loss(hi_res_true, hi_res_gen,
-                                          **calc_loss_kwargs)
-                loss, loss_details = loss_out
-
-                grad = tape.gradient(loss, training_weights)
-
+        with tf.device(device_name), tf.GradientTape(
+                watch_accessed_variables=False) as tape:
+            self.timer(tape.watch, training_weights)
+            hi_res_exo = self.timer(self.get_high_res_exo_input, hi_res_true)
+            hi_res_gen = self.timer(self._tf_generate, low_res, hi_res_exo)
+            loss_out = self.timer(self.calc_loss, hi_res_true, hi_res_gen,
+                                  **calc_loss_kwargs)
+            loss, loss_details = loss_out
+            grad = self.timer(tape.gradient, loss, training_weights)
         return grad, loss_details