Added support for reshape and flatten layers

kaffe/layers.py

@@ -22,7 +22,7 @@
     'EuclideanLoss': shape_scalar,
     'Eltwise': shape_identity,
     'Exp': shape_identity,
-    'Flatten': shape_not_implemented,
+    'Flatten': flatten_shape,
     'HDF5Data': shape_data,
     'HDF5Output': shape_identity,
     'HingeLoss': shape_scalar,
@@ -49,6 +49,7 @@
     'TanH': shape_identity,
     'WindowData': shape_not_implemented,
     'Threshold': shape_identity,
+    'Reshape':reshape_shape
 }
 
 LAYER_TYPES = LAYER_DESCRIPTORS.keys()

kaffe/shapes.py

@@ -24,7 +24,7 @@ def get_strided_kernel_output_shape(node, round_func):
 
 
 def shape_not_implemented(node):
-    raise NotImplementedError
+    raise NotImplementedError  
 
 
 def shape_identity(node):
@@ -69,7 +69,35 @@ def shape_concat(node):
             output_shape[axis] += parent.output_shape[axis]
     return tuple(output_shape)
 
-
+def reshape_shape(node) :
+
+    input_shape = node.get_only_parent().output_shape
+    input_shape_pr = input_shape.channels*input_shape.height*input_shape.width
+    input_shape_arr = [input_shape.batch_size,input_shape.channels,input_shape.height,input_shape.width]
+    pr = 1
+    axes = node.parameters.shape.dim
+    new_shape = [input_shape.batch_size,1,1,1]
+    for j in range(1,len(axes)) :
+        if axes[j] == 0 :
+            new_shape[j] = input_shape_arr[j]
+            pr *= new_shape[j]
+        elif not axes[j] == -1 :
+            new_shape[j] = int(axes[j])
+            pr *= new_shape[j]
+        elif axes[j] == -1 :
+            new_shape[j] = -1
+
+    for j in range(1,len(new_shape)) :
+        if new_shape[j] == -1 :
+            new_shape[j] = int(input_shape_pr/pr)
+
+    return TensorShape(new_shape[0],new_shape[1],new_shape[2],new_shape[3])                
+
+def flatten_shape(node) :
+    shape1 = node.get_only_parent().output_shape
+
+    return TensorShape(shape1.batch_size,shape1.channels*shape1.height*shape1.width,1,1)
+
 def shape_convolution(node):
     return get_strided_kernel_output_shape(node, math.floor)
 

kaffe/tensorflow/network.py

@@ -177,7 +177,7 @@ def lrn(self, input, radius, alpha, beta, name, bias=1.0):
 
     @layer
     def concat(self, inputs, axis, name):
-        return tf.concat(concat_dim=axis, values=inputs, name=name)
+        return tf.concat(axis=axis, values=inputs, name=name)
 
     @layer
     def add(self, inputs, name):
@@ -238,7 +238,24 @@ def batch_normalization(self, input, name, scale_offset=True, relu=False):
                 output = tf.nn.relu(output)
             return output
 
+
     @layer
     def dropout(self, input, keep_prob, name):
         keep = 1 - self.use_dropout + (self.use_dropout * keep_prob)
         return tf.nn.dropout(input, keep, name=name)
+
+    @layer
+    def reshape(self,input,b,x,y,c,name,transpose = False) :
+        if transpose :
+            input = tf.reshape(input,[-1,c,x,y])
+            return tf.transpose(input,(0,2,3,1))
+
+        return tf.reshape(input,[-1,x,y,c],name = name)
+
+    @layer
+    def flatten(self,input,name):
+        input = tf.transpose(input,(0,3,1,2))
+        dim = 1
+        for d in input.get_shape()[1:].as_list():
+                dim *= d
+        return tf.reshape(input,[-1,dim],name = name)

kaffe/tensorflow/transformer.py

@@ -139,14 +139,20 @@ def map_lrn(self, node):
         return TensorFlowNode('lrn', int(params.local_size / 2), alpha, params.beta)
 
     def map_concat(self, node):
-        axis = (2, 3, 1, 0)[node.parameters.axis]
+        if node.parents[0].kind == 'Flatten':
+            axis = node.parameters.axis
+        else :    
+            axis = (2, 3, 1, 0)[node.parameters.axis]
         return TensorFlowNode('concat', axis)
 
     def map_dropout(self, node):
         return TensorFlowNode('dropout', node.parameters.dropout_ratio)
 
     def map_batch_norm(self, node):
-        scale_offset = len(node.data) == 4
+        if node.data :
+            scale_offset = len(node.data) == 4
+        else :
+            scale_offset = True    
         kwargs = {} if scale_offset else {'scale_offset': False}
         return MaybeActivated(node, default=False)('batch_normalization', **kwargs)
 
@@ -158,6 +164,25 @@ def map_eltwise(self, node):
         except KeyError:
             raise KaffeError('Unknown elementwise operation: {}'.format(op_code))
 
+    def map_reshape(self,node) :
+
+        shape = node.output_shape
+        new_shape = [0]*4
+        new_shape[0] = shape[0]
+        new_shape[1] = shape[2]
+        new_shape[2] = shape[3]
+        new_shape[3] = shape[1]
+        parent_shape = node.get_only_parent().output_shape
+
+        ## we need to transpose again if a fc layer is reshaped to conv
+        kwargs = {'transpose' : False}
+        if parent_shape.height == 1 and parent_shape.width == 1 :
+            kwargs['transpose'] = True
+        return TensorFlowNode('reshape',new_shape[0],new_shape[1],new_shape[2],new_shape[3],**kwargs)
+
+    def map_flatten(self,node) :
+        return TensorFlowNode('flatten')
+
     def commit(self, chains):
         return chains