Added shape calculation logic for custom layers

README.md

@@ -291,7 +291,7 @@ Models                                                                      | Ca
 [Pix2Pix](https://github.com/phillipi/pix2pix)                              |   √   |   ×   |     ×      |
 [VQA](https://github.com/iamaaditya/VQA_Demo)                               |   √   |   √   |     √      |
 [Denoising Auto-Encoder](https://blog.keras.io/building-autoencoders-in-keras.html)                               |   ×   |   √   |     √      |
-[CapsNet](https://arxiv.org/abs/1710.09829)                               |   ×   |   √   |     √      |
+[CapsNet](https://arxiv.org/abs/1710.09829)                               |   ×   |   √   |     ×      |
 
 Note: For models that use a custom LRN layer (Alexnet), Keras expects the custom layer to be passed when it is loaded from json. LRN.py is located in keras_app/custom_layers. [Alexnet import for Keras](https://github.com/Cloud-CV/Fabrik/blob/master/tutorials/keras_custom_layer_usage.md)
 

ide/tasks.py

@@ -13,7 +13,7 @@
 from keras_app.views.layers_export import data, convolution, deconvolution, pooling, dense, dropout, embed,\
     recurrent, batch_norm, activation, flatten, reshape, eltwise, concat, upsample, locally_connected,\
     permute, repeat_vector, regularization, masking, gaussian_noise, gaussian_dropout, alpha_dropout, \
-    bidirectional, time_distributed, lrn, depthwiseConv
+    bidirectional, time_distributed, lrn, depthwiseConv, capsule_layer, squash, length, mask_capsule
 from keras_app.custom_layers import config as custom_layers_config
 from keras.models import model_from_json
 import tensorflow as tf
@@ -110,7 +110,11 @@ def export_keras_json(net, net_name, is_tf, reply_channel):
     }
 
     custom_layers_map = {
-        'LRN': lrn
+        'LRN': lrn,
+        'CapsuleLayer': capsule_layer,
+        'Length': length,
+        'MaskCapsule': mask_capsule,
+        'Squash': squash
     }
 
     # Remove any duplicate activation layers (timedistributed and bidirectional layers)
@@ -199,7 +203,6 @@ def isProcessPossible(layerId):
             })
         })
         return
-
     while(len(stack)):
         if ('Loss' in net[layerId]['info']['type'] or
                 net[layerId]['info']['type'] == 'Accuracy'):
@@ -213,6 +216,10 @@ def isProcessPossible(layerId):
             if (net[layerId]['info']['type'] != 'Scale'):
                 layer_in = [net_out[inputId]
                             for inputId in net[layerId]['connection']['input']]
+                # Handling MaskCapsule layer seperately
+                if (net[layerId]['info']['type'] == 'MaskCapsule'):
+                    layer_in = [[net_out[inputId]
+                                 for inputId in net[layerId]['connection']['input']]]
             # Need to check if next layer is Scale
             if (net[layerId]['info']['type'] == 'BatchNorm'):
                 idNext = net[layerId]['connection']['output'][0]

ide/utils/shapes.py

@@ -139,6 +139,15 @@ def repeat(layer):
     return shape
 
 
+def capsule(layer):
+    if (layer['info']['type'] == "CapsuleLayer"):
+        return [layer['params']['dim_capsule'], layer['params']['num_capsule']]
+    elif (layer['info']['type'] == "MaskCapsule"):
+        return [layer['shape']['input'][0][0] * layer['shape']['input'][0][1]]
+    else:
+        return [layer['shape']['input'][0]]
+
+
 def handle_concat_layer(outputLayer, inputLayer):
     if('input' not in outputLayer['shape']):
         shape = inputLayer['shape']['output'][:]
@@ -149,6 +158,17 @@ def handle_concat_layer(outputLayer, inputLayer):
     return shape
 
 
+def handle_mask_capsule_layer_input(inlayer, masklayer):
+    if inlayer['info']['type'] == "CapsuleLayer":
+        shape = inlayer['shape']['output'][:]
+        return shape
+    else:
+        shape = []
+        shape.append(masklayer['shape']['input'][:])
+        shape.append(inlayer['shape']['output'][:])
+        return shape
+
+
 def get_layer_shape(layer):
     # separating checking the type of layer inorder to make it modular
     # which can be reused in case we only want to get shapes of a single
@@ -161,6 +181,9 @@ def get_layer_shape(layer):
     elif(layer['info']['type'] in ['Convolution', 'Pooling', 'Deconvolution', 'DepthwiseConv']):
         return filter(layer)
 
+    elif(layer['info']['type'] in ['CapsuleLayer', 'Length', 'MaskCapsule']):
+        return capsule(layer)
+
     elif(layer['info']['type'] in ['InnerProduct', 'Recurrent', 'RNN', 'LSTM', 'Embed']):
         return output(layer)
 
@@ -208,7 +231,6 @@ def get_shapes(net):
     while(len(queue)):
         # using deque as stack
         layerId = queue.pop()
-
         if(net[layerId]['info']['type'] in dataLayers):
             net[layerId]['shape']['input'], net[layerId]['shape']['output'] = get_layer_shape(net[layerId])
         else:
@@ -219,6 +241,9 @@ def get_shapes(net):
                 # Handling Concat layer separately
                 if (net[outputId]['info']['type'] == "Concat"):
                     net[outputId]['shape']['input'] = handle_concat_layer(net[outputId], net[layerId])
+                elif (net[outputId]['info']['type'] == "MaskCapsule"):
+                    net[outputId]['shape']['input'] = handle_mask_capsule_layer_input(net[layerId],
+                                                                                      net[outputId])
                 else:
                     net[outputId]['shape']['input'] = net[layerId]['shape']['output'][:]
 

keras_app/custom_layers/capsule_layer.py

@@ -27,6 +27,9 @@ def build(self, input_shape):
         self.b = self.add_weight(shape=[input_shape[1], self.num_capsule],
                                  initializer=self.bias_initializer,
                                  name='b')
+        self.c = self.add_weight(shape=[input_shape[1], self.num_capsule],
+                                 initializer=self.bias_initializer,
+                                 name='c')
         super(CapsuleLayer, self).build(input_shape)
 
     def call(self, inputs, training=None):
@@ -36,6 +39,8 @@ def call(self, inputs, training=None):
         input_shape = K.shape(inputs_hat)
         b = self.b
         b = K.expand_dims(b, axis=0)
+        c = self.c
+        c = K.expand_dims(c, axis=0)
         assert self.num_routing > 0
         for i in range(self.num_routing):
             c = K.softmax(b)

keras_app/views/layers_export.py

@@ -136,6 +136,7 @@ def squash(layer, layer_in, layerId):
     axis = layer['params']['axis']
     out = {}
     out[layerId] = Squash(axis=axis)(*layer_in)
+    return out
 
 
 def dropout(layer, layer_in, layerId, tensor=True):