From 39074f22a75aac84c029e8ad9cb7873b5cdf4431 Mon Sep 17 00:00:00 2001 From: Piv <18462828+Piv200@users.noreply.github.com> Date: Sun, 21 Mar 2021 09:51:45 +1030 Subject: [PATCH] Add working train and eval functions for nyu_v2 --- fast_depth_functional.py | 120 +++++++++++++++++++++++++++++++-------- main.py | 14 ++--- 2 files changed, 101 insertions(+), 33 deletions(-) diff --git a/fast_depth_functional.py b/fast_depth_functional.py index b27aceb..70dd07a 100644 --- a/fast_depth_functional.py +++ b/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)) diff --git a/main.py b/main.py index 2bdb6ca..903e485 100644 --- a/main.py +++ b/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)