Fix VisionTransformer convert bug

examples/oneflow2onnx/models/GPU/test_vision_transformer.py

@@ -0,0 +1,49 @@
+"""
+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.nn as nn
+from typing import Union, List, Dict, Any, cast
+from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check
+
+from flowvision.models import ModelCreator
+
+import tempfile
+
+vit = ModelCreator.create_model("vit_base_patch16_224", pretrained=False)
+vit = vit.to("cuda")
+vit.eval()
+
+
+class VitGraph(flow.nn.Graph):
+    def __init__(self):
+        super().__init__()
+        self.m = vit
+
+    def build(self, x):
+        out = self.m(x)
+        return out
+
+
+def test_vit():
+
+    vit_graph = VitGraph()
+    vit_graph._compile(flow.randn(1, 3, 224, 224).to("cuda"))
+
+    convert_to_onnx_and_check(vit_graph, onnx_model_path=".", device="gpu")
+
+
+test_vit()

examples/oneflow2onnx/nodes/GPU/test_flatten_transpose.py

@@ -0,0 +1,55 @@
+"""
+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 tempfile
+import oneflow as flow
+from oneflow_onnx.oneflow2onnx.util import convert_to_onnx_and_check
+
+
+class FlattenTranspose(flow.nn.Module):
+    def __init__(self) -> None:
+        super(FlattenTranspose, self).__init__()
+
+    def forward(self, x: flow.Tensor) -> flow.Tensor:
+        res = x.flatten(2).transpose(1, 2)
+        print(res.shape)
+        return res
+
+
+flatten_transpose = FlattenTranspose()
+flatten_transpose = flatten_transpose.to("cuda")
+
+
+class FlattenTransposeOpGraph(flow.nn.Graph):
+    def __init__(self):
+        super().__init__()
+        self.m = flatten_transpose
+
+    def build(self, x):
+        out = self.m(x)
+        return out
+
+
+def test_flatten_transpose():
+
+    flatten_transpose_graph = FlattenTransposeOpGraph()
+    flatten_transpose_graph._compile(flow.randn(1, 3, 224, 224).to("cuda"))
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        flow.save(flatten_transpose_graph.state_dict(), tmpdirname)
+        convert_to_onnx_and_check(flatten_transpose_graph, onnx_model_path="/tmp", device="gpu")
+
+
+test_flatten_transpose()

oneflow_onnx/oneflow2onnx/handlers/array.py

@@ -109,35 +109,40 @@ def Version_5(cls, ctx, node, **kwargs):
             ctx.CopyShape(node.output_tensor_names[0], output_cast.output_tensor_names[0])
 
 
-@flow_op("flatten", "Flatten")
+@flow_op("flatten")
 class Flatten:
     @classmethod
     def Version_1(cls, ctx, node, **kwargs):
+        shape = ctx.get_shape(node.input_tensor_names[0])
+        dim = len(shape)
         start_dim = node.attrs.get("start_dim", 1)
-        dtype = ctx.get_dtype(node.input_tensor_names[0])
-        assert dtype == 1, f"onnx opset version 1/9 only support float32 data_type!"
-        assert start_dim >= 0, f"oneflow flatten can't support neagetive dim now!"
-        node.attrs["axis"] = start_dim
-
-    @classmethod
-    def Version_9(cls, ctx, node, **kwargs):
-        start_dim = node.attrs.get("start_dim", 1)
-        dtype = ctx.get_dtype(node.input_tensor_names[0])
-        assert dtype == 1, f"onnx opset version 1/9 only support float32 data_type!"
-        assert start_dim >= 0, f"oneflow flatten can't support neagetive dim now!"
-        node.attrs["axis"] = start_dim
-
-    @classmethod
-    def Version_11(cls, ctx, node, **kwargs):
-        start_dim = node.attrs.get("start_dim", 1)
-        assert start_dim >= 0, f"oneflow flatten can't support neagetive dim now!"
-        node.attrs["axis"] = start_dim
+        end_dim = node.attrs.get("end_dim", -1)
+        if end_dim < 0:
+            end_dim += dim
+        if start_dim == 1 and end_dim == dim - 1:
+            ctx.RemoveNode(node.name)
+            ctx.MakeNode("Flatten", [node.input_tensor_names[0]], attr={"aixs": start_dim}, outputs=[node.output_tensor_names[0]], op_name_scope=node.name, name="new_flatten")
+            return
+        if start_dim == 0 and end_dim == dim - 2:
+            ctx.RemoveNode(node.name)
+            ctx.MakeNode("Flatten", [node.input_tensor_names[0]], attr={"aixs": end_dim + 1}, outputs=[node.output_tensor_names[0]], op_name_scope=node.name, name="new_flatten")
+            return
 
-    @classmethod
-    def Version_13(cls, ctx, node, **kwargs):
-        start_dim = node.attrs.get("start_dim", 1)
-        assert start_dim >= 0, f"oneflow flatten can't support neagetive dim now!"
-        node.attrs["axis"] = start_dim
+        if start_dim > 1:
+            flatten_node = ctx.MakeNode("Flatten", [node.input_tensor_names[0]], attr={"aixs": 0}, op_name_scope=node.name, name="new_flatten")
+            new_shape = []
+            for i in range(start_dim):
+                new_shape.append(shape[i])
+            shape2 = 1
+            for i in range(start_dim, end_dim + 1):
+                shape2 *= shape[i]
+            new_shape.append(shape2)
+            for i in range(end_dim + 1, dim):
+                new_shape.append(shape[i])
+            ctx.RemoveNode(node.name)
+            new_shape_name = oneflow._oneflow_internal.UniqueStr("new_shape")
+            ctx.MakeConst(new_shape_name, np.array(new_shape, dtype=np.int64))
+            ctx.MakeNode("Reshape", [flatten_node.output_tensor_names[0], new_shape_name], outputs=[node.output_tensor_names[0]], op_name_scope=node.name, name="new_reshape")
 
 
 @flow_op("squeeze", "Squeeze")

oneflow_onnx/oneflow2onnx/util.py

@@ -146,7 +146,6 @@ def convert_to_onnx_and_check(
                 elif str(value.dtype) == "float" or str(value.dtype) == "float32":
                     value_tensor = flow.tensor(value, dtype=flow.float32, **device_kwargs)
                 elif str(value.dtype) == "float64":
-                    print("enter here")
                     value_tensor = flow.tensor(value, dtype=flow.float64, **device_kwargs)
                 elif str(value.dtype) == "bool":
                     value_tensor = flow.tensor(value, dtype=flow.bool, **device_kwargs)