Fix VisionTransformer convert bug

examples/oneflow2onnx/models/requirements.txt

@@ -1 +1 @@
-flowvision==0.0.3
+flowvision==0.2.0

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()

examples/oneflow2onnx/nodes/GPU/test_matmul.py

@@ -24,7 +24,7 @@ def __init__(self) -> None:
         self.matmul = flow.nn.Linear(20, 30)
 
     def forward(self, x: flow.Tensor) -> flow.Tensor:
-        return self.matmul(x)
+        return self.matmul(x) + flow.matmul(x, flow.ones(20, 1).to("cuda"), alpha=0.2)
 
 
 matmul = MatMul()

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/handlers/math.py

@@ -499,6 +499,7 @@ class MatMul:
     def Version_1(cls, ctx, node, **kwargs):
         transpose_a = node.attrs.get("transpose_a", 0)
         transpose_b = node.attrs.get("transpose_b", 0)
+        alpha = node.attrs.get("alpha")
 
         if transpose_a != 0:
             shape = ctx.get_shape(node.input_tensor_names[0])
@@ -524,6 +525,14 @@ def Version_1(cls, ctx, node, **kwargs):
             if val != 0:
                 raise ValueError(node.op_type + " attribute " + i + " is not supported")
 
+        if alpha != 1.0:
+            dtypes = node.output_dtypes
+            alpha = ctx.MakeConst(oneflow._oneflow_internal.UniqueStr("alpha"), np.array(alpha, dtype=util.Onnx2NumpyDtype(dtypes[0])))
+            mul = ctx.InsertNewNodeOnOutput("Mul", node.output_tensor_names[0], op_name_scope=node.name, name="mul_alpha")
+            mul.input_tensor_names.append(alpha.output_tensor_names[0])
+            ctx.set_dtype(mul.output_tensor_names[0], ctx.get_dtype(node.output_tensor_names[0]))
+            ctx.CopyShape(node.output_tensor_names[0], mul.output_tensor_names[0])
+
 
 @flow_op("erf", onnx_op="Erf")
 class Erf:

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)