add flatten op export onnx

docs/oneflow2onnx/op_list.md

@@ -22,4 +22,4 @@
 | 58 | ReduceSum| 59 | ReduceProd | 60 | ArgMax   | 61 | ArgMin             |
 |62  | Reshape  | 63 | Squeeze    | 64 | Transpose| 65 | Concat             |
 | 66 | Cast     | 67 | Identity   | 68 | Mul      | 69 | PReLU              |
-| 70 | LeakyReLU| 71 | Constant   |
+| 70 | LeakyReLU| 71 | Constant   | 72 | Flatten  |

examples/oneflow2onnx/nodes/test_flatten.py

+"""
+Copyright 2020 The OneFlow Authors. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import oneflow as flow
+import oneflow.typing as tp
+from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check
+
+
+def test_flatten():
+    @flow.global_function()
+    def flatten(x: tp.Numpy.Placeholder((3, 4, 2, 5))):
+        return flow.flatten(x, start_dim=1, end_dim=-1)
+
+    convert_to_onnx_and_check(flatten)
+

oneflow_onnx/oneflow2onnx/handlers/array.py

@@ -112,6 +112,13 @@ class Reshape:
                 node.output_tensor_names[0], output_cast.output_tensor_names[0]
             )
 
+@flow_op("flatten", "Flatten")
+class Flatten:
+    @classmethod
+    def Version_1(cls, ctx, node, **kwargs):
+        start_dim = node.attrs.get("start_dim", None)
+        node.attrs["axis"] = start_dim
+
 
 @flow_op("squeeze", "Squeeze")
 class Squeeze: