add const return support

paddle2onnx/mapper/exporter.cc

@@ -360,17 +360,46 @@ namespace paddle2onnx
 
         // 构建 else 分支图
         auto else_node_name = input_info[0].name;
-        auto conditional_block_cood_it = sub_block_map_.find(else_node_name);
-        Assert(conditional_block_cood_it != sub_block_map_.end(), "Don't find select_input else_input node.");
-        auto conditional_block_cood = conditional_block_cood_it->second;
-        auto else_graph = ExportConditionalBlock(parser, conditional_block_cood.first, conditional_block_cood.second, else_node_name);
-
+        ONNX_NAMESPACE::GraphProto else_graph;
+        if (else_node_name.find("fill_constant") != std::string::npos) {
+          else_graph.set_name("PaddlePaddle Const SubGraph " +
+                              std::to_string(block_id) + " " +
+                              std::to_string(op_id) + "_0");
+          for (auto i = 0; i < temp_helper.nodes.size(); ++i) {
+            auto node = temp_helper.nodes[i];
+            if (node->output(0) == else_node_name) {
+              else_graph.add_node()->CopyFrom(*node.get());
+              else_graph.add_output()->set_name(else_node_name);
+              break;
+            }
+          }
+        } else {
+          auto conditional_block_cood_it = sub_block_map_.find(else_node_name);
+          Assert(conditional_block_cood_it != sub_block_map_.end(), "Don't find select_input else_input node.");
+          auto conditional_block_cood = conditional_block_cood_it->second;
+          else_graph = ExportConditionalBlock(parser, conditional_block_cood.first, conditional_block_cood.second, else_node_name);
+        }
         // 构建 then 分支图
         auto then_node_name = input_info[1].name;
-        conditional_block_cood_it = sub_block_map_.find(then_node_name);
-        Assert(conditional_block_cood_it != sub_block_map_.end(), "Don't find select_input then_input node.");
-        conditional_block_cood = conditional_block_cood_it->second;
-        auto then_graph = ExportConditionalBlock(parser, conditional_block_cood.first, conditional_block_cood.second, then_node_name);
+        ONNX_NAMESPACE::GraphProto then_graph;
+        if (then_node_name.find("fill_constant") != std::string::npos) {
+          then_graph.set_name("PaddlePaddle Const SubGraph " +
+                              std::to_string(block_id) + " " +
+                              std::to_string(op_id) + "_1");
+          for (auto i = 0; i < temp_helper.nodes.size(); ++i) {
+            auto node = temp_helper.nodes[i];
+            if (node->output(0) == then_node_name) {
+              then_graph.add_node()->CopyFrom(*node.get());
+              then_graph.add_output()->set_name(then_node_name);
+              break;
+            }
+          }
+        } else {
+          auto conditional_block_cood_it = sub_block_map_.find(then_node_name);
+          Assert(conditional_block_cood_it != sub_block_map_.end(), "Don't find select_input then_input node.");
+          auto conditional_block_cood = conditional_block_cood_it->second;
+          then_graph = ExportConditionalBlock(parser, conditional_block_cood.first, conditional_block_cood.second, then_node_name);
+        }
 
         auto cond_info = parser.GetOpInput(block_id, op_id, "Mask");
         auto output_info = parser.GetOpOutput(block_id, op_id, "Out");

tests/onnxbase.py

@@ -64,7 +64,10 @@ def compare(result, expect, delta=1e-10, rtol=1e-10):
         # Convert Paddle Tensor to Numpy array
         if type(expect) == list:
             expect = expect[0]
-        expect = expect.numpy()
+        if isinstance(expect, paddle.Tensor):
+            expect = expect.numpy()
+        else:
+            expect = np.array(expect)
 
         # For result_shape is (1) and expect_shape shape is ()
         expect = expect.squeeze()

tests/test_ifelse.py

@@ -16,6 +16,7 @@
 from onnxbase import APIOnnx
 from onnxbase import randtool
 
+
 class BaseNet1(paddle.nn.Layer):
     def __init__(self):
         super(BaseNet1, self).__init__()
@@ -26,46 +27,80 @@ def forward(self, inputs):
         else:
             return inputs * 3
 
+
 def test_ifelse_1_true():
     op = BaseNet1()
     op.eval()
-    obj = APIOnnx(op, 'ifelse', [11])
+    obj = APIOnnx(op, "ifelse", [11])
     obj.set_input_data("input_data", paddle.to_tensor(1))
     obj.run()
 
+
 def test_ifelse_1_false():
     op = BaseNet1()
     op.eval()
-    obj = APIOnnx(op, 'ifelse', [11])
+    obj = APIOnnx(op, "ifelse", [11])
     obj.set_input_data("input_data", paddle.to_tensor(2))
     obj.run()
 
+
 class BaseNet2(paddle.nn.Layer):
     def __init__(self):
         super(BaseNet2, self).__init__()
 
     def forward(self, cond, inputs):
         if cond == 1:
-            return inputs * 1, inputs * 2 
+            return inputs * 1, inputs * 2
         else:
             return inputs * 3, inputs * 4
 
+
 def test_ifelse_2_true():
     op = BaseNet2()
     op.eval()
-    obj = APIOnnx(op, 'ifelse', [11])
+    obj = APIOnnx(op, "ifelse", [11])
     obj.set_input_data("input_data", paddle.to_tensor(1), paddle.to_tensor(1))
     obj.run()
 
+
 def test_ifelse_2_false():
     op = BaseNet2()
     op.eval()
-    obj = APIOnnx(op, 'ifelse', [11])
+    obj = APIOnnx(op, "ifelse", [11])
     obj.set_input_data("input_data", paddle.to_tensor(2), paddle.to_tensor(1))
     obj.run()
 
+
+class BaseNet3(paddle.nn.Layer):
+    def __init__(self):
+        super(BaseNet3, self).__init__()
+
+    def forward(self, inputs):
+        if inputs == 1:
+            return 2
+        else:
+            return 3
+
+
+def test_ifelse_3_true():
+    op = BaseNet3()
+    op.eval()
+    obj = APIOnnx(op, "ifelse", [11])
+    obj.set_input_data("input_data", paddle.to_tensor(1))
+    obj.run()
+
+
+def test_ifelse_3_false():
+    op = BaseNet3()
+    op.eval()
+    obj = APIOnnx(op, "ifelse", [11])
+    obj.set_input_data("input_data", paddle.to_tensor(2))
+    obj.run()
+
+
 if __name__ == "__main__":
     test_ifelse_1_true()
     test_ifelse_1_false()
     test_ifelse_2_true()
-    test_ifelse_2_false()
+    test_ifelse_2_false()
+    test_ifelse_3_false()