Merge pull request #852 from take-cheeze/tuple_input

[ONNX] Flatten input test data to support nested inputs

pytorch_pfn_extras/onnx/export_testcase.py

@@ -422,7 +422,8 @@ def write_to_pb(f: str, tensor: torch.Tensor, name: Optional[str] = None) -> Non
     os.makedirs(data_set_path, exist_ok=True)
     for pb_name in glob.glob(os.path.join(data_set_path, "*.pb")):
         os.remove(pb_name)
-    for i, (arg, name) in enumerate(zip(named_args, input_names)):
+    flat_inputs = torch._C._jit_flatten(named_args)[0]
+    for i, (arg, name) in enumerate(zip(flat_inputs, input_names)):
         f = os.path.join(data_set_path, 'input_{}.pb'.format(i))
         write_to_pb(f, arg, name)
 

tests/pytorch_pfn_extras_tests/onnx_tests/test_export_testcase.py

@@ -633,3 +633,31 @@ def custom(model, args, f, **kwargs):
     actual = ort_session.run(None, {"x": x.cpu().numpy()})[0]
     expected = model(x)
     np.testing.assert_allclose(actual, expected.detach().numpy(), atol=1e-3)
+
+
+def test_export_tuple_input():
+
+    class Net(nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.linear = nn.Linear(5, 10, bias=False)
+            self.in_channels = [5, 10]
+
+        def forward(self, inputs):
+            assert isinstance(inputs, tuple), f"{type(inputs)=}"
+            linears = [self.linear(x) for x in inputs]
+            return linears
+
+
+    model = Net()
+    x = torch.rand(2, 5)
+
+    export_testcase(
+        model,
+        ((x,),),
+        output_dir,
+        input_names=["x"],
+        training=model.training,
+        do_constant_folding=False,
+        opset_version=12,
+    )