Add working train and eval functions for nyu_v2

fast_depth_functional.py

@@ -1,5 +1,23 @@
 import tensorflow as tf
 import tensorflow.keras as keras
+import tensorflow_datasets as tfds
+
+
+# Ripped from: https://forums.developer.nvidia.com/t/could-not-create-cudnn-handle-cudnn-status-alloc-failed/108261/4?u=mpivato4
+# Seems to be an issue on windows so explicitly set gpu growth
+def fix_windows_gpu():
+    gpus = tf.config.experimental.list_physical_devices('GPU')
+    if gpus:
+        try:
+            # Currently, memory growth needs to be the same across GPUs
+            for gpu in gpus:
+                tf.config.experimental.set_memory_growth(gpu, True)
+            logical_gpus = tf.config.experimental.list_logical_devices('GPU')
+            print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
+        except RuntimeError as e:
+            # Memory growth must be set before GPUs have been initialized
+            print(e)
+
 
 '''
 Functional version of fastdepth model
@@ -17,11 +35,10 @@ def FDDepthwiseBlock(inputs,
     return keras.layers.ReLU(6., name='conv_pw_%d_relu' % block_id)(x)
 
 
-def make_fastdepth_functional(weights=None):
-    # This doesn't work, at least right now...
-    input = keras.layers.Input(shape=(224, 224, 3))
+def make_mobilenet_nnconv5(weights=None, shape=(224, 224, 3)):
+    input = keras.layers.Input(shape=shape)
     x = input
-    mobilenet = keras.applications.MobileNet(include_top=False, weights=weights)
+    mobilenet = keras.applications.MobileNet(input_tensor=x, include_top=False, weights=weights)
     for layer in mobilenet.layers:
         x = layer(x)
         if layer.name == 'conv_pw_5_relu':
@@ -55,43 +72,98 @@ def make_fastdepth_functional(weights=None):
     return keras.Model(inputs=input, outputs=x, name="fast_depth")
 
 
+# TODO: Fix these, float doesn't work same as pytorch
 def delta1_metric(y_true, y_pred):
     maxRatio = tf.maximum(y_pred / y_true, y_true / y_pred)
-    return float((maxRatio < 1.25).float().mean())
+    return tf.nn.moments(tf.cast(maxRatio < tf.convert_to_tensor(1.25), tf.float32), axes=None)[0]
 
 
 def delta2_metric(y_true, y_pred):
     maxRatio = tf.maximum(y_pred / y_true, y_true / y_pred)
-    return float((maxRatio < 1.2 ** 25).float().mean())
+    return tf.nn.moments(tf.cast(maxRatio < tf.convert_to_tensor(1.25 ** 2), tf.float32), axes=None)[0]
 
 
 def delta3_metric(y_true, y_pred):
     maxRatio = tf.maximum(y_pred / y_true, y_true / y_pred)
-    return float((maxRatio < 1.25 ** 3).float().mean())
+    return tf.nn.moments(tf.cast(maxRatio < tf.convert_to_tensor(1.25 ** 3), tf.float32), axes=None)[0]
 
 
-def fastdepth_for_training():
-    # Pretrained mobilenet on imagenet dataset
-    model = make_fastdepth_functional('imagenet')
-    return model.compile(optimizer=keras.optimizers.SGD(momentum=0.9),
-                         loss=keras.losses.MSE(),
-                         metrics=[keras.metrics.RootMeanSquaredError(),
-                                  keras.metrics.MeanSquaredError(),
-                                  delta1_metric,
-                                  delta2_metric,
-                                  delta3_metric])
+def compile(model):
+    # TODO: Learning rate (exponential decay)
+    model.compile(optimizer=keras.optimizers.SGD(momentum=0.9),
+                  loss=keras.losses.MeanSquaredError(),
+                  metrics=[keras.metrics.RootMeanSquaredError(),
+                           keras.metrics.MeanSquaredError(),
+                           delta1_metric,
+                           delta2_metric,
+                           delta3_metric])
 
 
-def train_compiled_model(compiled_model, dataset):
+def train(existing_model=None, pretrained_weights='imagenet', epochs=4, save_file=None, dataset=None):
+    if not existing_model:
+        existing_model = make_mobilenet_nnconv5(pretrained_weights)
+    compile(existing_model)
+    if not dataset:
+        dataset = load_nyu()
+    existing_model.fit(dataset, epochs=epochs)
+    if save_file:
+        existing_model.save(save_file)
+    return existing_model
+
+
+def evaluate(compiled_model, dataset=None):
     """
+    Evaluate the model using rmse, delta1/2/3 metrics
+    :param compiled_model: Compiled, trained model to evaluate
+    :param dataset: Dataset for evaluation. Should be of format {'image': image, 'depth': label},
+            where label width/height matches image width/height.
+            Defaults to Tensorflow nyu_v2 evaluation split dataset (https://www.tensorflow.org/datasets/catalog/nyu_depth_v2)
+    """
+    if not dataset:
+        dataset = load_nyu_evaluate()
+    compiled_model.evaluate(dataset, verbose=1)
 
-    :param compiled_model: Compiled model to train on
-    :param dataset: Dataset to train on (must be compatible with model
+
+def forward(model, image):
+    """
+    Propagate a single or batch of images through the model. Image(s) should be in format NHWC
+    :param model:
+    :param image:
     :return:
     """
-    # TODO: Use tf nyu_v2 dataset to train.
-    pass
+    return model(crop_and_resize(image))
 
 
-if __name__ == '__main__':
-    make_fastdepth_functional().summary()
+def load_model(file):
+    return keras.models.load_model(file, custom_objects={'delta1_metric': delta1_metric,
+                                                         'delta2_metric': delta2_metric,
+                                                         'delta3_metric': delta3_metric})
+
+
+def crop_and_resize(x):
+    shape = tf.shape(x['depth'])
+
+    def layer():
+        return keras.Sequential([
+            keras.layers.experimental.preprocessing.CenterCrop(shape[1], shape[2]),
+            keras.layers.experimental.preprocessing.Resizing(224, 224, interpolation='nearest')
+        ])
+
+    # Reshape label to 4d, can't use array unwrap as it's unsupported by tensorflow
+    return layer()(x['image']), layer()(tf.reshape(x['depth'], [shape[0], shape[1], shape[2], 1]))
+
+
+def load_nyu():
+    builder = tfds.builder('nyu_depth_v2')
+    builder.download_and_prepare(download_dir='../nyu')
+    return builder \
+        .as_dataset(split='train', shuffle_files=True) \
+        .shuffle(buffer_size=1024) \
+        .batch(8) \
+        .map(lambda x: crop_and_resize(x))
+
+
+def load_nyu_evaluate():
+    builder = tfds.builder('nyu_depth_v2')
+    builder.download_and_prepare(download_dir='../nyu')
+    return builder.as_dataset(split='validation').batch(1).map(lambda x: crop_and_resize(x))

main.py

@@ -1,11 +1,7 @@
-import tensorflow_datasets as tfds
-
-
-def load_nyu():
-    builder = tfds.builder('nyu_depth_v2')
-    builder.download_and_prepare(download_dir='../nyu')
-    return builder.as_dataset(split='train', shuffle_files=True)
-
+import fast_depth_functional as fd
 
 if __name__ == '__main__':
-    load_nyu()
+    fd.fix_windows_gpu()
+    model = fd.load_model('fast_depth_nyu_v2_224_224_3_e1')
+    fd.compile(model)
+    fd.evaluate(model)