From 625ecba731b8e93a3055612d3d68f3220b0aa15e Mon Sep 17 00:00:00 2001 From: Michael Pivato Date: Sun, 8 Aug 2021 09:25:55 +0000 Subject: [PATCH] Add small option to packnet, fix docs and first/final conv layers (prev 32) --- packnet_functional.py | 55 ++++++++++++++++++++++++++----------------- 1 file changed, 33 insertions(+), 22 deletions(-) diff --git a/packnet_functional.py b/packnet_functional.py index 3929df6..f2f8765 100644 --- a/packnet_functional.py +++ b/packnet_functional.py @@ -52,7 +52,7 @@ def packnet_inverse_depth(inputs, out_channels=1, min_depth=0.5): def pack_3d(inputs, kernel_size, r=2, features_3d=8): """ - Implementatino of the 3d packing block proposed here: https://arxiv.org/abs/1905.02693 + Implementation of the 3d packing block proposed here: https://arxiv.org/abs/1905.02693 :param inputs: Tensor inputs :param kernel_size: Conv3D kernels size :param r: Packing factor @@ -79,17 +79,22 @@ def unpack_3d(inputs, out_channels, kernel_size, r=2, features_3d=8): # TODO: Support different channel format (right now we're supporting NHWC, we should also support NCHW) -def make_packnet(shape=(224, 224, 3), skip_add=True, features_3d=4, dropout=None): +def make_packnet(shape=(224, 224, 3), skip_add=False, features_3d=8, dropout=None, small=False): """ Make the PackNet depth network. :param shape: Input shape of the image :param skip_add: Set to use add rather than concat skip connections, defaults to True - :return: + :param features_3d: Number of layers in 3D conv for packing/unpacking layers + :param dropout: Whether to build the model with dropout layers fpr regularisation. Useful in training only + :param small: Set True to not include the middle-most layer. Reduces params from ~128M -> ~34M + Further reductions can be achieved by using additive skip connections and less 3d features (down to min ~10M) + :return: Packnet Keras Model """ # ================ ENCODER ================= input = keras.layers.Input(shape=shape) - x = packnet_conv2d(input, 32, 5, 1) + initial_conv_channels = 32 if small else 64 + x = packnet_conv2d(input, initial_conv_channels, 5, 1) skip_1 = x x = packnet_conv2d(x, 64, 7, 1) x = pack_3d(x, 5, features_3d=features_3d) @@ -103,53 +108,59 @@ def make_packnet(shape=(224, 224, 3), skip_add=True, features_3d=4, dropout=None x = residual_block(x, 256, 3, 1, dropout) x = pack_3d(x, 3, features_3d=features_3d) skip_5 = x - x = residual_block(x, 512, 3, 1, dropout) - x = pack_3d(x, 3, features_3d=features_3d) + if not small: + x = residual_block(x, 512, 3, 1, dropout) + x = pack_3d(x, 3, features_3d=features_3d) # ================ ENCODER ================= # ================ DECODER ================= # Addition requires we half the outputs so there is a matching number of channels divide_factor = (2 if skip_add else 1) - # layer 12 - 13 - x = unpack_3d(x, 512 // divide_factor, 3, features_3d=features_3d) - x = keras.layers.Add()( - [x, skip_5]) if skip_add else keras.layers.Concatenate()([x, skip_5]) - x = packnet_conv2d(x, 512, 3, 1) - # layer 14 - 15 + # layer 7 + if not small: + x = unpack_3d(x, 512 // divide_factor, 3, features_3d=features_3d) + x = keras.layers.Add()( + [x, skip_5]) if skip_add else keras.layers.Concatenate()([x, skip_5]) + x = packnet_conv2d(x, 512, 3, 1) + # layer 8 x = unpack_3d(x, 256 // divide_factor, 3, features_3d=features_3d) x = keras.layers.Add()( [x, skip_4]) if skip_add else keras.layers.Concatenate()([x, skip_4]) x = packnet_conv2d(x, 256, 3, 1) layer_8 = x - # layer 16 + # layer 9 x = packnet_inverse_depth(x, 1) - # layer 17 - 18 + # layer 10 u_layer_8 = unpack_3d(layer_8, 128 // divide_factor, 3, features_3d=features_3d) x = keras.layers.UpSampling2D()(x) x = keras.layers.Add()([u_layer_8, skip_3, x]) if skip_add else keras.layers.Concatenate()([u_layer_8, skip_3, x]) x = packnet_conv2d(x, 128, 3, 1) layer_10 = x - # layer 19 + # layer 11 x = packnet_inverse_depth(x, 1) - # layer 20 - 21 + # layer 12 u_layer_10 = unpack_3d(layer_10, 64, 3, features_3d=features_3d) x = keras.layers.UpSampling2D()(x) x = keras.layers.Add()([u_layer_10, skip_2, x]) if skip_add else keras.layers.Concatenate()([u_layer_10, skip_2, x]) x = packnet_conv2d(x, 64, 3, 1) layer_12 = x - # layer 22 + # layer 13 x = packnet_inverse_depth(x) - # layer 23 - 24 - u_layer_12 = unpack_3d(layer_12, 64, 3, features_3d=features_3d) + # layer 14 + u_layer_12 = unpack_3d(layer_12, initial_conv_channels, 3, features_3d=features_3d) x = keras.layers.UpSampling2D()(x) x = keras.layers.Add()([u_layer_12, skip_1, x]) if skip_add else keras.layers.Concatenate()([u_layer_12, skip_1, x]) - x = packnet_conv2d(x, 64, 3, 1) - # layer 25 + x = packnet_conv2d(x, initial_conv_channels, 3, 1) + # layer 15 x = packnet_inverse_depth(x) # ================ DECODER ================= return keras.Model(inputs=input, outputs=x, name="PackNet") + if __name__ == '__main__': # This is the implementation used by the packnet sfm paper - make_packnet(features_3d=8, skip_add=False).summary() \ No newline at end of file + make_packnet().summary() + + # This is the very small version of packnet + make_packnet(small=True, features_3d=1, skip_add=True).summary()