diff --git a/examples/oneflow2onnx/models/requirements.txt b/examples/oneflow2onnx/models/requirements.txt index 115c052..079be6c 100644 --- a/examples/oneflow2onnx/models/requirements.txt +++ b/examples/oneflow2onnx/models/requirements.txt @@ -1 +1 @@ -flowvision==0.0.3 +flowvision==0.2.0 diff --git a/examples/oneflow2onnx/nodes/GPU/test_flatten_transpose.py b/examples/oneflow2onnx/nodes/GPU/test_flatten_transpose.py new file mode 100644 index 0000000..bdf08b9 --- /dev/null +++ b/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() diff --git a/examples/oneflow2onnx/nodes/GPU/test_matmul.py b/examples/oneflow2onnx/nodes/GPU/test_matmul.py index 103d8f1..19c536c 100644 --- a/examples/oneflow2onnx/nodes/GPU/test_matmul.py +++ b/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() diff --git a/oneflow_onnx/oneflow2onnx/handlers/array.py b/oneflow_onnx/oneflow2onnx/handlers/array.py index 6cc87e7..a22d49f 100644 --- a/oneflow_onnx/oneflow2onnx/handlers/array.py +++ b/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") diff --git a/oneflow_onnx/oneflow2onnx/handlers/math.py b/oneflow_onnx/oneflow2onnx/handlers/math.py index 8aa2b20..8ff3fa2 100644 --- a/oneflow_onnx/oneflow2onnx/handlers/math.py +++ b/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: diff --git a/oneflow_onnx/oneflow2onnx/util.py b/oneflow_onnx/oneflow2onnx/util.py index 6257398..19a848a 100644 --- a/oneflow_onnx/oneflow2onnx/util.py +++ b/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)