Add small option to packnet, fix docs and first/final conv layers (prev 32)

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)
+    if not small:
-    x = pack_3d(x, 3, features_3d=features_3d)
+        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
+    # layer 7
-    x = unpack_3d(x, 512 // divide_factor, 3, features_3d=features_3d)
+    if not small:
-    x = keras.layers.Add()(
+        x = unpack_3d(x, 512 // divide_factor, 3, features_3d=features_3d)
-        [x, skip_5]) if skip_add else keras.layers.Concatenate()([x, skip_5])
+        x = keras.layers.Add()(
-    x = packnet_conv2d(x, 512, 3, 1)
+            [x, skip_5]) if skip_add else keras.layers.Concatenate()([x, skip_5])
-    # layer 14 - 15
+        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
+    # layer 14
-    u_layer_12 = unpack_3d(layer_12, 64, 3, features_3d=features_3d)
+    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)
+    x = packnet_conv2d(x, initial_conv_channels, 3, 1)
-    # layer 25
+    # 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()
+    make_packnet().summary()
+
+    # This is the very small version of packnet
+    make_packnet(small=True, features_3d=1, skip_add=True).summary()