Flatten Support for ONNX(#98)

ROCm · Apr 12, 2019 · f05574d · f05574d
1 parent f55c266
commit f05574d
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 2 deletions.
diff --git a/model_compiler/python/nnir.py b/model_compiler/python/nnir.py
@@ -384,6 +384,8 @@ def updateLocals(self):
                 elif node.type in ['reshape']:
                     input = self.tensor_dict[node.inputs[0]]
                     param = node.attr.get('shape')
+                    if not param:
+                        param = input.shape
                     icount = 1
                     ocount = 1
                     out_shape = [0,0,0,0]                    
@@ -468,8 +470,12 @@ def updateLocals(self):
                 elif node.type in ['flatten']:
                     input = self.tensor_dict[node.inputs[0]]
                     axis = node.attr.get("axis")
-                    if axis == 1:
+                    if axis == 0:
+                        shape = [1, input.shape[0]*input.shape[1]*input.shape[2]*input.shape[3], 1, 1]
+                    elif axis == 1:
                         shape = [input.shape[0], input.shape[1]*input.shape[2]*input.shape[3], 1, 1]
+                    else:
+                        raise ValueError("Flatten: unsupoorted flatten: " + str(axis))
                     local = IrTensor()
                     local.setName(output)
                     local.setInfo(input.type, shape)

diff --git a/model_compiler/python/onnx_to_nnir.py b/model_compiler/python/onnx_to_nnir.py
@@ -63,7 +63,8 @@
     'GlobalAveragePool'  : 'global_avg_pool',
     'Softmax'            : 'softmax',
     'Reshape'            : 'reshape',
-    'Transpose'          : 'transpose'
+    'Transpose'          : 'transpose',
+    'Flatten'            : 'flatten'
 }
 
 onnx2ir_data_type = [