From 8def90cf6c0435cf1578024193c61f77a08bbcf6 Mon Sep 17 00:00:00 2001
From: Christophe Ecabert <christophe.ecabert@gmail.com>
Date: Thu, 15 Aug 2019 18:50:48 +0200
Subject: [PATCH] CHG: Add Reshape + flatten from
 https://github.com/ethereon/caffe-tensorflow/pull/147

---
 kaffe/layers.py                 |  7 ++++---
 kaffe/shapes.py                 | 30 ++++++++++++++++++++++++++++++
 kaffe/tensorflow/network.py     | 20 ++++++++++++++++++--
 kaffe/tensorflow/transformer.py | 25 +++++++++++++++++++++++--
 4 files changed, 75 insertions(+), 7 deletions(-)

diff --git a/kaffe/layers.py b/kaffe/layers.py
index b3d9fb2..9a4fed1 100644
--- a/kaffe/layers.py
+++ b/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,
@@ -38,7 +38,9 @@
     'Normalize': shape_identity,
     'Pooling': shape_pool,
     'Power': shape_identity,
+    'PReLU': shape_identity,
     'ReLU': shape_identity,
+    'Reshape': reshape_shape,
     'Scale': shape_identity,
     'Sigmoid': shape_identity,
     'SigmoidCrossEntropyLoss': shape_scalar,
@@ -49,8 +51,7 @@
     'Slice': shape_not_implemented,
     'TanH': shape_identity,
     'WindowData': shape_not_implemented,
-    'Threshold': shape_identity,
-    'PReLU': shape_identity,
+    'Threshold': shape_identity
 }
 
 LAYER_TYPES = LAYER_DESCRIPTORS.keys()
diff --git a/kaffe/shapes.py b/kaffe/shapes.py
index a70ff14..0a384bf 100644
--- a/kaffe/shapes.py
+++ b/kaffe/shapes.py
@@ -70,6 +70,36 @@ def shape_concat(node):
     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)
 
diff --git a/kaffe/tensorflow/network.py b/kaffe/tensorflow/network.py
index 6f36c56..4ec4f4a 100644
--- a/kaffe/tensorflow/network.py
+++ b/kaffe/tensorflow/network.py
@@ -255,11 +255,27 @@ def dropout(self, input, keep_prob, name):
         return tf.nn.dropout(input, keep, name=name)
         
     @layer
-    def l2_normalize(self, input):
+    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)
+        
+    @layer
+    def l2_normalize(self, input, name):
         # NOTE: Currently, only inference is supported
         with tf.variable_scope(name) as scope:
           shp = input.get_shape().as_list()
           outputs = tf.nn.l2_normalize(x=input, axis=-1)
-          alpha = self.make_var('alpha', shape=[-1:])
+          alpha = self.make_var('alpha', shape=shp[-1:])
           outputs = tf.multiply(outputs, alpha)
           return outputs
diff --git a/kaffe/tensorflow/transformer.py b/kaffe/tensorflow/transformer.py
index 9564fe5..4a9de10 100644
--- a/kaffe/tensorflow/transformer.py
+++ b/kaffe/tensorflow/transformer.py
@@ -142,7 +142,10 @@ 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):
@@ -161,9 +164,27 @@ 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 map_normalize(self, node):
         return TensorFlowNode('l2_normalize')
-
+        
     def commit(self, chains):
         return chains