fix unsqueeze bug

examples/oneflow2onnx/models/GPU/test_mobilenet_v3.py

@@ -44,4 +44,14 @@ def test_mobilenetv3():
         flow.save(mobilenetv3.state_dict(), tmpdirname)
         convert_to_onnx_and_check(mobilenetv3_graph, flow_weight_dir=tmpdirname, onnx_model_path=".", device="gpu")
 
+def test_mobilenetv3_opset14():
+
+    mobilenetv3_graph = MobileNetV3()
+    mobilenetv3_graph._compile(flow.randn(1, 3, 224, 224).to("cuda"))
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        flow.save(mobilenetv3.state_dict(), tmpdirname)
+        convert_to_onnx_and_check(mobilenetv3_graph, flow_weight_dir=tmpdirname, onnx_model_path=".", device="gpu", opset=14)
+
 test_mobilenetv3()
+test_mobilenetv3_opset14()

examples/oneflow2onnx/nodes/CPU/test_conv.py

@@ -45,4 +45,15 @@ def test_conv2d():
         flow.save(conv_module.state_dict(), tmpdirname)
         convert_to_onnx_and_check(conv_graph, flow_weight_dir=tmpdirname, onnx_model_path="/tmp")
 
+def test_conv2d_opset14():
+
+    conv_graph = Conv2dOpGraph()
+    conv_graph._compile(flow.randn(1, 3, 224, 224))
+
+    with tempfile.TemporaryDirectory() as tmpdirname:
+        flow.save(conv_module.state_dict(), tmpdirname)
+        convert_to_onnx_and_check(conv_graph, flow_weight_dir=tmpdirname, onnx_model_path="/tmp", opset=14)
+
+
 test_conv2d()
+test_conv2d_opset14()

oneflow_onnx/oneflow2onnx/handlers/math.py

@@ -90,6 +90,22 @@ def Version_6(cls, ctx, node, **kwargs):
         ctx.set_dtype(node.input_tensor_names[1], unsqueeze_dtype)
         super().Version_6(ctx, node, **kwargs)
 
+    @classmethod
+    def Version_13(cls, ctx, node, **kwargs):
+        axis = node.attrs["axis"]
+        unsqueeze_axes = []
+        x_rank = len(ctx.get_shape(node.input_tensor_names[0]))
+        for i in range(x_rank):
+            if axis != i:
+                unsqueeze_axes.append(i)
+        assert len(ctx.get_shape(node.input_tensor_names[1])) == 1
+        shape_node = ctx.MakeConst(
+            oneflow._oneflow_internal.UniqueStr("shape"), np.array(unsqueeze_axes)
+        )
+        ctx.InsertNewNodeOnInput(
+            node, "Unsqueeze", [node.input_tensor_names[1], shape_node.output_tensor_names[0]]
+        )
+        super().Version_6(ctx, node, **kwargs)
 
 @flow_op(["leaky_relu", "softplus"], onnx_op=["LeakyRelu", "Softplus"])
 class DirectOpSinceOpset1:
@@ -280,6 +296,10 @@ def Version_1(cls, ctx, node, **kwargs):
         ctx.MakeNode(
             "Mul", [node.input_tensor_names[0], node1.output_tensor_names[0]], outputs=[node.output_tensor_names[0]], op_name_scope=node.name, name="mul"
         )
+
+    @classmethod
+    def Version_14(cls, ctx, node, **kwargs):
+        pass
 
 @flow_op("hardsigmoid", onnx_op="HardSigmoid")
 class HardSigmoid:

oneflow_onnx/oneflow2onnx/optimizer/const_fold_optimizer.py

@@ -144,7 +144,10 @@ def _FoldUnsqueeze(node, graph):
         numpy expand_dims only supports to unsqueeze one dim one time, so reshape is used to simplify the logic
         """
         const_val = node.input_nodes[0].get_tensor_value(as_list=False)
-        axes = node.attrs["axes"]
+        if "axes" in node.attrs:
+            axes = node.attrs["axes"]
+        else:
+            axes = node.input_nodes[1].get_tensor_value(as_list=False)
         util.MakeSure(
             all(axis >= 0 for axis in axes),
             "onnx spec says it only supports positive axis",