support var op convert

examples/oneflow2onnx/nodes/CPU/test_var.py

@@ -0,0 +1,45 @@
+"""
+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 Var(flow.nn.Module):
+    def __init__(self) -> None:
+        super(Var, self).__init__()
+
+    def forward(self, x: flow.Tensor) -> flow.Tensor:
+        y = flow.var(x, dim=None, unbiased=True, keepdim=True)
+        return y
+
+var_module = Var()
+class VarOpGraph(flow.nn.Graph):
+    def __init__(self):
+        super().__init__()
+        self.m = var_module
+
+    def build(self, x):
+        out = self.m(x)
+        return out
+
+def test_var():
+
+    var_op_graph = VarOpGraph()
+    var_op_graph._compile(flow.arange(48, dtype=flow.float32).reshape(2, 2, 3, 4))
+    convert_to_onnx_and_check(var_op_graph, onnx_model_path="/tmp", opset=13)
+
+test_var()
+

examples/oneflow2onnx/nodes/GPU/test_var.py

@@ -0,0 +1,46 @@
+"""
+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 Var(flow.nn.Module):
+    def __init__(self) -> None:
+        super(Var, self).__init__()
+
+    def forward(self, x: flow.Tensor) -> flow.Tensor:
+        y = flow.var(x, dim=None, unbiased=True, keepdim=True)
+        return y
+
+var_module = Var()
+var_module = var_module.to("cuda")
+class VarOpGraph(flow.nn.Graph):
+    def __init__(self):
+        super().__init__()
+        self.m = var_module
+
+    def build(self, x):
+        out = self.m(x)
+        return out
+
+def test_var():
+
+    var_op_graph = VarOpGraph()
+    var_op_graph._compile(flow.arange(48, dtype=flow.float32).reshape(2, 2, 3, 4).to("cuda"))
+    convert_to_onnx_and_check(var_op_graph, onnx_model_path="/tmp", opset=13, device="gpu")
+
+test_var()
+

oneflow_onnx/oneflow2onnx/handlers/math.py

@@ -302,7 +302,7 @@ def Version_14(cls, ctx, node, **kwargs):
         pass
 
 @flow_op("silu", onnx_op="Mul")
-class HardSwish:
+class Silu:
     @classmethod
     def Version_1(cls, ctx, node, **kwargs):
         dtypes = node.output_dtypes
@@ -929,3 +929,63 @@ def Version_7(cls, ctx, node, **kwargs):
         ctx.CopyShape(output_name, new_node.output_tensor_names[0])
         ctx.set_dtype(new_node.output_tensor_names[0], ctx.get_dtype(output_name))
 
+
+@flow_op(["var"])
+class Var:
+    @classmethod
+    def Version_13(cls, ctx, node, **kwargs):
+        origin_dim = node.attrs.get("dim", None)
+        unbiased = node.attrs.get("unbiased", None)
+        keepdim = node.attrs.get("keepdim", None)
+        num_elements = 1
+        dtypes = node.output_dtypes
+        input_shape = ctx.get_shape(node.input_tensor_names[0])
+        keepdim_mean = 0 if origin_dim is None else keepdim
+
+        if origin_dim is None:
+            dim = []            
+            for i in range(len(input_shape)):
+                num_elements *= input_shape[i]
+                dim.append(i)
+            reduce_mean_node = ctx.MakeNode(
+                "ReduceMean", [node.input_tensor_names[0]], op_name_scope=node.name, name="reduce_mean", dtypes=dtypes, attr={"axes":dim, "keepdims": 0}
+            )
+            t_mean = reduce_mean_node.output_tensor_names[0]
+
+        else:
+            reduce_mean_node = ctx.MakeNode(
+                "ReduceMean", [node.input_tensor_names[0]], op_name_scope=node.name, name="reduce_mean", dtypes=dtypes, attr={"axes":origin_dim, "keepdims": 1}
+            )
+            t_mean = reduce_mean_node.output_tensor_names[0]
+            for i in range(len(origin_dim)):
+                num_elements *= input_shape[i]
+
+        sub_node = ctx.MakeNode(
+            "Sub", [node.input_tensor_names[0], t_mean], op_name_scope=node.name, name="sub", dtypes=dtypes
+        )
+        sub_v = sub_node.output_tensor_names[0]
+        mul_node = ctx.MakeNode(
+            "Mul", [sub_v, sub_v], op_name_scope=node.name, name="mul", dtypes=dtypes
+        )
+        sqr_sub = mul_node.output_tensor_names[0]
+        if unbiased is None:
+            unbiased = False
+
+        ctx.RemoveNode(node.name)
+        if unbiased:
+            var_node = ctx.MakeNode(
+                "ReduceMean", [sqr_sub], op_name_scope=node.name, name="var", dtypes=dtypes, attr={"axes":origin_dim, "keepdims": keepdim_mean}
+            )
+            var = var_node.output_tensor_names[0]
+            scalar_node = ctx.MakeConst(
+                oneflow._oneflow_internal.UniqueStr("scalar"), np.array([num_elements]).astype(np.float32)
+            )
+            one = ctx.MakeConst(oneflow._oneflow_internal.UniqueStr("constant"), np.array([1]).astype(np.float32))
+            num_elements = scalar_node.output_tensor_names[0]
+            mul = ctx.MakeNode("Mul", [var, num_elements])
+            sub = ctx.MakeNode("Sub", [num_elements, one.output_tensor_names[0]])
+            var = ctx.MakeNode("Div", [mul.output_tensor_names[0], sub.output_tensor_names[0]], outputs=[node.output_tensor_names[0]])
+        else:
+            var_node = ctx.MakeNode(
+                "ReduceMean", [sqr_sub], op_name_scope=node.name, name="var", dtypes=dtypes, attr={"axes":origin_dim, "keepdims": keepdim_mean}, outputs=[node.output_tensor_names[0]]
+            )