Merge pull request #146 from nasa/develop

* Support mixed precision
* USGS script updates
* Minor bugfixes

Author: bcoltin

delta/config/delta.yaml

@@ -79,6 +79,7 @@ train:
   callbacks:       ~
   # augmentation functions to apply to training data
   augmentations:   ~
+  disable_mixed_precision: false
   validation:
     # if true, skips the first steps from the training set to use for validation instead
     from_training: true

delta/extensions/losses.py

@@ -199,7 +199,7 @@ def preprocess(self, y_true, y_pred):
             nodata = tf.expand_dims(nodata, -1)
 
         # zero all nodata entries
-        y_pred = y_pred * tf.cast(tf.logical_not(nodata), tf.float32)
+        y_pred = tf.cast(y_pred, tf.float32) * tf.cast(tf.logical_not(nodata), tf.float32)
 
         true_convert = tf.cast(tf.logical_not(nodata), tf.float32) * true_convert
         return (true_convert, y_pred)

delta/extensions/sources/sentinel1.py

@@ -178,6 +178,8 @@ def _prep(self, paths):
         assert isinstance(paths, str)
         ext = os.path.splitext(paths)[1]
 
+        self._name = os.path.splitext(os.path.basename(paths))[0]
+
         tif_path = None
         if ext == '.zip': # Need to unpack
 

delta/extensions/sources/tiff.py

@@ -310,8 +310,6 @@ def callback_function(output_roi, data, _):
             block_x = output_roi.min_x // ts[1]
             block_y = output_roi.min_y // ts[0]
 
-            print('Write output_roi = ' + str(output_roi))
-
             # Loop on bands
             if len(data.shape) == 2:
                 writer.write_block(data[:, :], block_y, block_x, 0)

delta/imagery/utilities.py

@@ -69,6 +69,6 @@ def progress_bar(text, fill_amount, prefix = '', length = 80): #pylint: disable=
         Number of characters to fill as bar
     """
     filled_length = int(length * fill_amount)
-    fill_char = '█' if sys.stdout.encoding.lower() == 'utf-8' else 'X'
+    fill_char = '█' if str(sys.stdout.encoding).lower() == 'utf-8' else 'X'
     prog_bar = fill_char * filled_length + '-' * (length - filled_length)
     print('\r%s |%s| %s' % (prefix, prog_bar, text), end = '\r')

delta/ml/ml_config.py

@@ -180,9 +180,15 @@ def __init__(self):
         self.register_field('steps', int, None, config.validate_non_negative, 'Batches to train per epoch.')
         self.register_field('optimizer', (str, dict), None, None, 'Keras optimizer to use.')
         self.register_field('callbacks', list, 'callbacks', None, 'Callbacks used to modify training')
+        self.register_field('disable_mixed_precision', bool, 'disable_mixed_precision', None,
+                            'Disables mixed precision tensorflow policy. By default DELTA will use mixed '
+                            'precision if the hardware supports it. Details on ways to improve mixed '
+                            'precision performance: '
+                            'https://www.tensorflow.org/guide/mixed_precision#summary')
         self.register_arg('epochs', '--epochs')
         self.register_arg('batch_size', '--batch-size')
         self.register_arg('steps', '--steps')
+        self.register_arg('disable_mixed_precision', '--disable-mixed-precision', action="store_true", type=None)
         self.register_field('augmentations', list, None, None, None)
         self.register_component(ValidationConfig(), 'validation')
         self.register_component(NetworkConfig(), 'network')

delta/subcommands/train.py

@@ -22,20 +22,54 @@
 import sys
 import time
 
-#import logging
-#logging.getLogger("tensorflow").setLevel(logging.DEBUG)
+# import logging
+# logging.getLogger("tensorflow").setLevel(logging.DEBUG)
 
 import tensorflow as tf
+from tensorflow.keras import mixed_precision
 
 from delta.config import config
 from delta.imagery import imagery_dataset
 from delta.ml.train import train
 from delta.ml.config_parser import config_model
 from delta.ml.io import save_model, load_model
 
-#tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.DEBUG)
+
+# tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.DEBUG)
+
+def mixed_policy_device_compatible():
+
+    # gpu check logic taken from https://github.com/keras-team/keras/blob/70d7d07bd186b929d81f7a8ceafff5d78d8bd701/keras/mixed_precision/device_compatibility_check.py # pylint: disable=line-too-long
+    gpus = tf.config.list_physical_devices('GPU')
+    gpu_details_list = [tf.config.experimental.get_device_details(g) for g in gpus]
+
+    supported_device_strs = []
+    unsupported_device_strs = []
+    for details in gpu_details_list:
+        name = details.get('device_name', 'Unknown GPU')
+        cc = details.get('compute_capability')
+        if cc:
+            device_str = '%s, compute capability %s.%s' % (name, cc[0], cc[1])
+            if cc >= (7, 0):
+                supported_device_strs.append(device_str)
+            else:
+                unsupported_device_strs.append(device_str)
+        else:
+            unsupported_device_strs.append(
+                name + ', no compute capability (probably not an Nvidia GPU)')
+
+    if unsupported_device_strs or not supported_device_strs:
+        return False
+    # else mixed policy is compatible
+    return True
+
 
 def main(options):
+    if mixed_policy_device_compatible() and not config.train.disable_mixed_precision():
+        mixed_precision.set_global_policy('mixed_float16')
+        print('Tensorflow Mixed Precision is enabled. This improves training performance on compatible GPUs. '
+              'However certain precautions should be taken and several additional changes can be made to improve '
+              'performance further. Details: https://www.tensorflow.org/guide/mixed_precision#summary')
 
     images = config.dataset.images()
     if not images: