Merge pull request #178 from NREL/gb/dev

Significant refactor on handling of feature names in DataHandler and model

sup3r/models/base.py

@@ -72,20 +72,21 @@ def __init__(self,
             history
         meta : dict | None
             Model meta data that describes how the model was created.
-        means : np.ndarray | list | None
-            Set of mean values for data normalization with the same length as
-            number of features. Can be used to maintain a consistent
-            normalization scheme between transfer learning domains.
-        stdevs : np.ndarray | list | None
-            Set of stdev values for data normalization with the same length as
-            number of features. Can be used to maintain a consistent
-            normalization scheme between transfer learning domains.
+        means : dict | None
+            Set of mean values for data normalization keyed by feature name.
+            Can be used to maintain a consistent normalization scheme between
+            transfer learning domains.
+        stdevs : dict | None
+            Set of stdev values for data normalization keyed by feature name.
+            Can be used to maintain a consistent normalization scheme between
+            transfer learning domains.
         default_device : str | None
             Option for default device placement of model weights. If None and a
             single GPU exists, that GPU will be the default device. If None and
             multiple GPUs exist, the CPU will be the default device (this was
             tested as most efficient given the custom multi-gpu strategy
-            developed in self.run_gradient_descent())
+            developed in self.run_gradient_descent()). Examples: "/gpu:0" or
+            "/cpu:0"
         name : str | None
             Optional name for the GAN.
         """
@@ -116,10 +117,8 @@ def __init__(self,
         self._gen = self.load_network(gen_layers, 'generator')
         self._disc = self.load_network(disc_layers, 'discriminator')
 
-        self._means = (means if means is None else np.array(means).astype(
-            np.float32))
-        self._stdevs = (stdevs if stdevs is None else np.array(stdevs).astype(
-            np.float32))
+        self._means = means
+        self._stdevs = stdevs
 
     def save(self, out_dir):
         """Save the GAN with its sub-networks to a directory.
@@ -224,8 +223,12 @@ def discriminate(self, hi_res, norm_in=False):
             hi_res = hi_res.numpy()
 
         if norm_in and self._means is not None:
+            mean_arr = [self._means[fn] for fn in self.hr_out_features]
+            std_arr = [self._stdevs[fn] for fn in self.hr_out_features]
+            mean_arr = np.array(mean_arr, dtype=np.float32)
+            std_arr = np.array(std_arr, dtype=np.float32)
             hi_res = hi_res if isinstance(hi_res, tf.Tensor) else hi_res.copy()
-            hi_res = (hi_res - self._means) / self._stdevs
+            hi_res = (hi_res - mean_arr) / std_arr
 
         out = self.discriminator.layers[0](hi_res)
         for i, layer in enumerate(self.discriminator.layers[1:]):
@@ -324,10 +327,6 @@ def model_params(self):
         -------
         dict
         """
-        means = (self._means
-                 if self._means is None else [float(m) for m in self._means])
-        stdevs = (self._stdevs if self._stdevs is None else
-                  [float(s) for s in self._stdevs])
 
         config_optm_g = self.get_optimizer_config(self.optimizer)
         config_optm_d = self.get_optimizer_config(self.optimizer_disc)
@@ -338,8 +337,8 @@ def model_params(self):
             'version_record': self.version_record,
             'optimizer': config_optm_g,
             'optimizer_disc': config_optm_d,
-            'means': means,
-            'stdevs': stdevs,
+            'means': self._means,
+            'stdevs': self._stdevs,
             'meta': self.meta,
         }
 
@@ -383,7 +382,7 @@ def init_weights(self, lr_shape, hr_shape, device=None):
 
         with tf.device(device):
             hr_exo_data = {}
-            for feature in self.exogenous_features:
+            for feature in self.hr_exo_features:
                 hr_exo_data[feature] = hr_exo
             _ = self._tf_generate(low_res, hr_exo_data)
             _ = self._tf_discriminate(hi_res)
@@ -514,6 +513,7 @@ def calc_loss_disc(disc_out_true, disc_out_gen):
 
         return loss_disc
 
+    @tf.function
     def calc_loss(self,
                   hi_res_true,
                   hi_res_gen,
@@ -645,7 +645,8 @@ def train_epoch(self,
             present, each batch from the batch_handler will be divided up
             between the GPUs and the resulting gradient from each GPU will
             constitute a single gradient descent step with the nominal learning
-            rate that the model was initialized with.
+            rate that the model was initialized with. If true and multiple gpus
+            are found, default_device device should be set to /gpu:0
 
         Returns
         -------
@@ -845,7 +846,8 @@ def train(self,
             present, each batch from the batch_handler will be divided up
             between the GPUs and the resulting gradient from each GPU will
             constitute a single gradient descent step with the nominal learning
-            rate that the model was initialized with.
+            rate that the model was initialized with. If true and multiple gpus
+            are found, default_device device should be set to /gpu:0
         """
 
         self.set_norm_stats(batch_handler.means, batch_handler.stds)
@@ -854,9 +856,9 @@ def train(self,
             s_enhance=batch_handler.s_enhance,
             t_enhance=batch_handler.t_enhance,
             smoothing=batch_handler.smoothing,
-            training_features=batch_handler.training_features,
-            train_only_features=batch_handler.train_only_features,
-            output_features=batch_handler.output_features,
+            lr_features=batch_handler.lr_features,
+            hr_exo_features=batch_handler.hr_exo_features,
+            hr_out_features=batch_handler.hr_out_features,
             smoothed_features=batch_handler.smoothed_features)
 
         epochs = list(range(n_epoch))

sup3r/models/conditional_moments.py

@@ -45,14 +45,14 @@ def __init__(self, gen_layers,
             history
         meta : dict | None
             Model meta data that describes how the model was created.
-        means : np.ndarray | list | None
-            Set of mean values for data normalization with the same length as
-            number of features. Can be used to maintain a consistent
-            normalization scheme between transfer learning domains.
-        stdevs : np.ndarray | list | None
-            Set of stdev values for data normalization with the same length as
-            number of features. Can be used to maintain a consistent
-            normalization scheme between transfer learning domains.
+        means : dict | None
+            Set of mean values for data normalization keyed by feature name.
+            Can be used to maintain a consistent normalization scheme between
+            transfer learning domains.
+        stdevs : dict | None
+            Set of stdev values for data normalization keyed by feature name.
+            Can be used to maintain a consistent normalization scheme between
+            transfer learning domains.
         default_device : str | None
             Option for default device placement of model weights. If None and a
             single GPU exists, that GPU will be the default device. If None and
@@ -84,10 +84,8 @@ def __init__(self, gen_layers,
 
         self._gen = self.load_network(gen_layers, 'generator')
 
-        self._means = (means if means is None
-                       else np.array(means).astype(np.float32))
-        self._stdevs = (stdevs if stdevs is None
-                        else np.array(stdevs).astype(np.float32))
+        self._means = means
+        self._stdevs = stdevs
 
     def save(self, out_dir):
         """Save the model with its sub-networks to a directory.
@@ -174,10 +172,6 @@ def model_params(self):
         -------
         dict
         """
-        means = (self._means if self._means is None
-                 else [float(m) for m in self._means])
-        stdevs = (self._stdevs if self._stdevs is None
-                  else [float(s) for s in self._stdevs])
 
         config_optm_g = self.get_optimizer_config(self.optimizer)
 
@@ -189,8 +183,8 @@ def model_params(self):
                         'num_par': num_par,
                         'version_record': self.version_record,
                         'optimizer': config_optm_g,
-                        'means': means,
-                        'stdevs': stdevs,
+                        'means': self._means,
+                        'stdevs': self._stdevs,
                         'meta': self.meta,
                         }
 
@@ -395,9 +389,9 @@ def train(self, batch_handler,
             s_enhance=batch_handler.s_enhance,
             t_enhance=batch_handler.t_enhance,
             smoothing=batch_handler.smoothing,
-            train_only_features=batch_handler.train_only_features,
-            training_features=batch_handler.training_features,
-            output_features=batch_handler.output_features,
+            lr_features=batch_handler.lr_features,
+            hr_exo_features=batch_handler.hr_exo_features,
+            hr_out_features=batch_handler.hr_out_features,
             smoothed_features=batch_handler.smoothed_features)
 
         epochs = list(range(n_epoch))

sup3r/models/linear.py

@@ -17,11 +17,11 @@ class LinearInterp(AbstractInterface):
     """Simple model to do linear interpolation on the spatial and temporal axes
     """
 
-    def __init__(self, features, s_enhance, t_enhance, t_centered=False):
+    def __init__(self, lr_features, s_enhance, t_enhance, t_centered=False):
         """
         Parameters
         ----------
-        features : list
+        lr_features : list
             List of feature names that this model will operate on for both
             input and output. This must match the feature axis ordering in the
             array input to generate().
@@ -35,7 +35,7 @@ def __init__(self, features, s_enhance, t_enhance, t_centered=False):
             time-centered (e.g. interp 01:00 02:00 to 00:45 01:15 01:45 02:15)
         """
 
-        self._features = features
+        self._lr_features = lr_features
         self._s_enhance = s_enhance
         self._t_enhance = t_enhance
         self._t_centered = t_centered
@@ -78,27 +78,31 @@ class init args.
     @property
     def meta(self):
         """Get meta data dictionary that defines the model params"""
-        return {'features': self._features,
+        return {'lr_features': self._lr_features,
                 's_enhance': self._s_enhance,
                 't_enhance': self._t_enhance,
                 't_centered': self._t_centered,
-                'training_features': self.training_features,
-                'output_features': self.output_features,
+                'hr_out_features': self.hr_out_features,
                 'class': self.__class__.__name__,
                 }
 
     @property
-    def training_features(self):
+    def lr_features(self):
         """Get the list of input feature names that the generative model was
         trained on.
         """
-        return self._features
+        return self._lr_features
 
     @property
-    def output_features(self):
+    def hr_out_features(self):
         """Get the list of output feature names that the generative model
         outputs"""
-        return self._features
+        return self._lr_features
+
+    @property
+    def hr_exo_features(self):
+        """Returns an empty list for LinearInterp model"""
+        return []
 
     def save(self, out_dir):
         """
@@ -142,7 +146,7 @@ def generate(self, low_res, norm_in=False, un_norm_out=False,
                     int(low_res.shape[1] * self._s_enhance),
                     int(low_res.shape[2] * self._s_enhance),
                     int(low_res.shape[3] * self._t_enhance),
-                    len(self.output_features))
+                    len(self.hr_out_features))
         logger.debug('LinearInterp model with s_enhance of {} '
                      'and t_enhance of {} '
                      'downscaling low-res shape {} to high-res shape {}'