Fixed double-flatten issue.

Author: gmalivenko

pytorch2keras/layers.py

@@ -232,14 +232,19 @@ def convert_flatten(params, w_name, scope_name, inputs, layers, weights, short_n
         short_names: use short names for keras layers
     """
     print('Conerting reshape ...')
+
     if short_names:
         tf_name = 'R' + random_string(7)
     else:
         tf_name = w_name + str(random.random())
 
     # TODO: check if the input is already flattened
-    reshape = keras.layers.Flatten(name=tf_name)
-    layers[scope_name] = reshape(layers[inputs[0]])
+    # Ad-hoc to avoid it:
+    if len(list(layers[inputs[0]].shape)) == 2:
+        layers[scope_name] = layers[inputs[0]]
+    else:
+        reshape = keras.layers.Flatten(name=tf_name)
+        layers[scope_name] = reshape(layers[inputs[0]])
 
 
 def convert_gemm(params, w_name, scope_name, inputs, layers, weights, short_names):

tests/view.py

@@ -15,7 +15,7 @@ def __init__(self, inp=10, out=16, kernel_size=3, bias=True):
 
     def forward(self, x):
         x = self.conv2d(x)
-        x = x.view([x.size(0), -1, 2, 1, 1, 1, 1, 1])
+        x = x.view([x.size(0), -1, 2, 1, 1, 1, 1, 1]).view(x.size(0), -1).view(x.size(0), -1)
         x = torch.nn.Tanh()(x)
         return x