Gb/optm state

sup3r/models/abstract.py

@@ -1037,7 +1037,7 @@ def get_optimizer_config(optimizer):
         Parameters
         ----------
         optimizer : tf.keras.optimizers.Optimizer
-            TF-Keras optimizer object
+            TF-Keras optimizer object (e.g., Adam)
 
         Returns
         -------
@@ -1053,6 +1053,29 @@ def get_optimizer_config(optimizer):
                 conf[k] = int(v)
         return conf
 
+    @classmethod
+    def get_optimizer_state(cls, optimizer):
+        """Get a set of state variables for the optimizer
+
+        Parameters
+        ----------
+        optimizer : tf.keras.optimizers.Optimizer
+            TF-Keras optimizer object (e.g., Adam)
+
+        Returns
+        -------
+        state : dict
+            Optimizer state variables
+        """
+        lr = cls.get_optimizer_config(optimizer)['learning_rate']
+        state = {'learning_rate': lr}
+        for var in optimizer.variables:
+            name = var.name
+            var = var.numpy().flatten()
+            var = np.abs(var).mean()  # collapse ndarrays into mean absolute
+            state[name] = float(var)
+        return state
+
     @staticmethod
     def update_loss_details(loss_details, new_data, batch_len, prefix=None):
         """Update a dictionary of loss_details with loss information from a new

sup3r/models/base.py

@@ -746,8 +746,10 @@ def train_epoch(
 
             b_loss_details['gen_trained_frac'] = float(trained_gen)
             b_loss_details['disc_trained_frac'] = float(trained_disc)
+
             self.dict_to_tensorboard(b_loss_details)
             self.dict_to_tensorboard(self.timer.log)
+
             loss_details = self.update_loss_details(
                 loss_details,
                 b_loss_details,
@@ -1000,10 +1002,9 @@ def train(
                 loss_details['train_loss_gen'], loss_details['train_loss_disc']
             )
 
-            if all(
-                loss in loss_details
-                for loss in ('val_loss_gen', 'val_loss_disc')
-            ):
+            check1 = 'val_loss_gen' in loss_details
+            check2 = 'val_loss_disc' in loss_details
+            if check1 and check2:
                 msg += 'gen/disc val loss: {:.2e}/{:.2e} '.format(
                     loss_details['val_loss_gen'], loss_details['val_loss_disc']
                 )
@@ -1016,14 +1017,15 @@ def train(
                 'weight_gen_advers': weight_gen_advers,
                 'disc_loss_bound_0': disc_loss_bounds[0],
                 'disc_loss_bound_1': disc_loss_bounds[1],
-                'learning_rate_gen': self.get_optimizer_config(self.optimizer)[
-                    'learning_rate'
-                ],
-                'learning_rate_disc': self.get_optimizer_config(
-                    self.optimizer_disc
-                )['learning_rate'],
             }
 
+            opt_g = self.get_optimizer_state(self.optimizer)
+            opt_d = self.get_optimizer_state(self.optimizer_disc)
+            opt_g = {f'OptmGen/{key}': val for key, val in opt_g.items()}
+            opt_d = {f'OptmDisc/{key}': val for key, val in opt_d.items()}
+            extras.update(opt_g)
+            extras.update(opt_d)
+
             weight_gen_advers = self.update_adversarial_weights(
                 loss_details,
                 adaptive_update_fraction,