Merge branch 'master' into model_opsets

Author: asi1024

.flexci/config.pbtxt

@@ -77,6 +77,22 @@ configs {
   }
 }
 
+configs {
+  key: "pytorch-pfn-extras.torch112-linux"
+  value {
+    requirement {
+      cpu: 6
+      gpu: 2
+      memory: 32
+      disk: 10
+    }
+    time_limit: {
+      seconds: 1200
+    }
+    command: "bash .flexci/linux/script.sh torch112"
+  }
+}
+
 configs {
   key: "pytorch-pfn-extras.torch18-win"
   value {
@@ -144,3 +160,20 @@ configs {
     command: ".flexci\\windows\\run.bat torch111"
   }
 }
+
+configs {
+  key: "pytorch-pfn-extras.torch112-win"
+  value {
+    requirement {
+      cpu: 4
+      gpu: 2
+      memory: 24
+      disk: 10
+      image: "windows"
+    }
+    time_limit: {
+      seconds: 1200
+    }
+    command: ".flexci\\windows\\run.bat torch112"
+  }
+}

.flexci/linux/build_and_push.sh

@@ -8,7 +8,7 @@ if [ "${IMAGE_BASE}" = "" ]; then
 fi
 
 TEST_PIP_PACKAGES="
-matplotlib tensorboard ipython ipywidgets pandas optuna onnx onnxruntime
+matplotlib tensorboard ipython ipywidgets pandas optuna onnxruntime
 pytest flake8 pysen[lint] pytest-cov
 "
 
@@ -34,7 +34,7 @@ docker_build_and_push torch18 \
     --build-arg base_image="nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04" \
     --build-arg python_version="3.6.9" \
     --build-arg pip_install_torch_args="torch==1.8.* torchvision==0.9.* -f https://download.pytorch.org/whl/cu102/torch_stable.html" \
-    --build-arg pip_install_dep_args="cupy-cuda102 pytorch-ignite ${TEST_PIP_PACKAGES}" \
+    --build-arg pip_install_dep_args="cupy-cuda102 pytorch-ignite onnx==1.11.0 ${TEST_PIP_PACKAGES}" \
     &
 WAIT_PIDS="$! ${WAIT_PIDS}"
 
@@ -43,7 +43,7 @@ docker_build_and_push torch19 \
     --build-arg base_image="nvidia/cuda:10.2-cudnn7-devel-ubuntu18.04" \
     --build-arg python_version="3.9.5" \
     --build-arg pip_install_torch_args="torch==1.9.* torchvision==0.10.* -f https://download.pytorch.org/whl/cu102/torch_stable.html" \
-    --build-arg pip_install_dep_args="cupy-cuda102 pytorch-ignite ${TEST_PIP_PACKAGES}" \
+    --build-arg pip_install_dep_args="cupy-cuda102 pytorch-ignite onnx ${TEST_PIP_PACKAGES}" \
     &
 WAIT_PIDS="$! ${WAIT_PIDS}"
 
@@ -52,16 +52,31 @@ docker_build_and_push torch110 \
     --build-arg base_image="nvidia/cuda:11.3.1-cudnn8-devel-ubuntu18.04" \
     --build-arg python_version="3.9.7" \
     --build-arg pip_install_torch_args="torch==1.10.* torchvision==0.11.* -f https://download.pytorch.org/whl/cu113/torch_stable.html" \
-    --build-arg pip_install_dep_args="cupy-cuda113 pytorch-ignite ${TEST_PIP_PACKAGES}" \
+    --build-arg pip_install_dep_args="cupy-cuda113 pytorch-ignite onnx ${TEST_PIP_PACKAGES}" \
     &
 WAIT_PIDS="$! ${WAIT_PIDS}"
 
+# Wait until the build complete.
+for P in ${WAIT_PIDS}; do
+    wait ${P}
+done
+WAIT_PIDS=""
+
 # PyTorch 1.11 + Python 3.9
 docker_build_and_push torch111 \
     --build-arg base_image="nvidia/cuda:11.3.1-cudnn8-devel-ubuntu18.04" \
     --build-arg python_version="3.9.7" \
     --build-arg pip_install_torch_args="torch==1.11.* torchvision==0.12.* -f https://download.pytorch.org/whl/cu113/torch_stable.html" \
-    --build-arg pip_install_dep_args="cupy-cuda113 pytorch-ignite ${TEST_PIP_PACKAGES}" \
+    --build-arg pip_install_dep_args="cupy-cuda113 pytorch-ignite onnx ${TEST_PIP_PACKAGES}" \
+    &
+WAIT_PIDS="$! ${WAIT_PIDS}"
+
+# PyTorch 1.12 + Python 3.9
+docker_build_and_push torch112 \
+    --build-arg base_image="nvidia/cuda:11.3.1-cudnn8-devel-ubuntu18.04" \
+    --build-arg python_version="3.9.7" \
+    --build-arg pip_install_torch_args="torch==1.12.* torchvision==0.13.* -f https://download.pytorch.org/whl/cu113/torch_stable.html" \
+    --build-arg pip_install_dep_args="cupy-cuda113 pytorch-ignite onnx ${TEST_PIP_PACKAGES}" \
     &
 WAIT_PIDS="$! ${WAIT_PIDS}"
 

.flexci/windows/test.ps1

@@ -37,6 +37,13 @@ if ($test -eq "torch18") {
     RunOrDie python -m pip install -U pip "setuptools<59.6"
     RunOrDieWithRetry 3 python -m pip install torch==1.11.* torchvision==0.12.* -f https://download.pytorch.org/whl/cu113/torch_stable.html
 
+} elseif ($test -eq "torch112") {
+    # PyTorch 1.12 + Python 3.9
+    ActivateCUDA 11.3
+    ActivatePython 3.9
+    RunOrDie python -m pip install -U pip "setuptools<59.6"
+    RunOrDieWithRetry 3 python -m pip install torch==1.12.* torchvision==0.13.* -f https://download.pytorch.org/whl/cu113/torch_stable.html
+
 } else {
     throw "Unsupported test variant: $test"
 }

.github/CODEOWNERS

@@ -1,3 +1,3 @@
 * @asi1024 @emcastillo @kmaehashi
-/pytorch_pfn_extras/onnx @take-cheeze @okdshin
-/tests/pytorch_pfn_extras_tests/onnx @take-cheeze @okdshin
+/pytorch_pfn_extras/onnx @xuzijian629 @okdshin
+/tests/pytorch_pfn_extras_tests/onnx @xuzijian629 @okdshin

.github/workflows/pretest.yml

@@ -7,7 +7,7 @@ jobs:
     runs-on: ubuntu-18.04
     strategy:
       matrix:
-        torch: ['1.8.*', '1.9.*', '1.10.*', '1.11.*']
+        torch: ['1.8.*', '1.9.*', '1.10.*', '1.11.*', '1.12.*']
 
     steps:
     - name: Checkout

pytorch_pfn_extras/onnx/_constants.py

@@ -0,0 +1,11 @@
+import pytorch_pfn_extras
+import torch.onnx
+import torch.onnx.symbolic_helper
+
+if pytorch_pfn_extras.requires("1.12.0"):
+    from torch.onnx._constants import *  # NOQA
+else:
+    onnx_default_opset = torch.onnx.symbolic_helper._default_onnx_opset_version
+    onnx_main_opset = torch.onnx.symbolic_helper._onnx_main_opset
+    onnx_stable_opsets = torch.onnx.symbolic_helper._onnx_stable_opsets
+    onnx_constant_folding_opsets = torch.onnx.symbolic_helper._constant_folding_opset_versions if pytorch_pfn_extras.requires("1.11.0") else torch.onnx.constant_folding_opset_versions  # type: ignore[attr-defined]

pytorch_pfn_extras/onnx/_globals.py

@@ -0,0 +1,29 @@
+import pytorch_pfn_extras
+import torch
+import torch.onnx.symbolic_helper
+from typing import Optional
+
+
+class _InternalGlobalsBeforeTorch1_11:
+    @property
+    def export_onnx_opset_version(self) -> int:
+        return torch.onnx.symbolic_helper._export_onnx_opset_version
+
+    @property
+    def operator_export_type(self) -> Optional[torch._C._onnx.OperatorExportTypes]:
+        return torch.onnx.symbolic_helper._operator_export_type
+
+    @property
+    def training_mode(self) -> Optional[torch._C._onnx.TrainingMode]:
+        return torch.onnx.symbolic_helper._training_mode
+
+    @property
+    def onnx_shape_inference(self) -> bool:
+        return torch.onnx.symbolic_helper._onnx_shape_inference
+
+
+if pytorch_pfn_extras.requires("1.12.0"):
+    import torch.onnx._globals
+    GLOBALS = torch.onnx._globals.GLOBALS
+else:
+    GLOBALS = _InternalGlobalsBeforeTorch1_11()

pytorch_pfn_extras/onnx/export_testcase.py

@@ -11,6 +11,7 @@
 import onnx
 import onnx.numpy_helper
 import pytorch_pfn_extras
+import pytorch_pfn_extras.onnx._constants
 import torch
 import torch.autograd
 from torch.onnx import OperatorExportTypes
@@ -151,12 +152,7 @@ def _export(
     bytesio = io.BytesIO()
     opset_ver = kwargs.get('opset_version', None)
     if opset_ver is None:
-        if pytorch_pfn_extras.requires('1.12.0'):
-            from torch.onnx._constants import onnx_default_opset
-            opset_ver = onnx_default_opset
-        else:
-            from torch.onnx.symbolic_helper import _default_onnx_opset_version
-            opset_ver = _default_onnx_opset_version
+        opset_ver = pytorch_pfn_extras.onnx._constants.onnx_default_opset
         kwargs['opset_version'] = opset_ver
     if use_pfto or not pytorch_pfn_extras.requires('1.10.0'):
         strip_doc_string = kwargs.get('strip_doc_string', True)

pytorch_pfn_extras/onnx/pfto_exporter/export.py

@@ -9,6 +9,8 @@
 import onnx.numpy_helper
 import onnx.shape_inference
 import pytorch_pfn_extras
+import pytorch_pfn_extras.onnx._constants
+from pytorch_pfn_extras.onnx._globals import GLOBALS
 import torch
 import torch.jit
 import torch.onnx.symbolic_helper as sym_hel
@@ -316,11 +318,7 @@ def optimize_onnx(self, graph: torch._C.Graph) -> torch._C.Graph:
         else:
             self.run_jit_pass(torch._C._jit_pass_onnx_scalar_type_analysis, graph)
 
-        opset_versions = (
-            sym_hel._constant_folding_opset_versions   # type: ignore[attr-defined]
-            if pytorch_pfn_extras.requires("1.11.0")
-            else torch.onnx.constant_folding_opset_versions)  # type: ignore[attr-defined]
-        if self.do_constant_folding and self.opset_version in opset_versions:
+        if self.do_constant_folding and self.opset_version in pytorch_pfn_extras.onnx._constants.onnx_constant_folding_opsets:
             folded: Dict[str, torch.IValue] = torch._C._jit_pass_onnx_constant_fold(  # type: ignore[attr-defined]
                 graph, self.vars, self.opset_version
             )
@@ -502,6 +500,8 @@ def run_symbolic_function(self, g: torch._C.Graph, n: torch._C.Node, sym_func: C
         node_inputs = list(n.inputs())
         if n.kind() == "prim::PythonOp":
             node_inputs.extend(n.scalar_args())
+            if "module" in attrs:
+                del attrs["module"]
         sym_outs = _to_tuple_if_not_sequence(sym_func(g, *node_inputs, **attrs))
         assert len(sym_outs) == n.outputsSize(), f"{sym_outs}: {len(sym_outs)} vs {n.outputsSize()}"
 
@@ -829,6 +829,8 @@ def apply_dynamic_axes_info(out: onnx.ValueInfoProto, k: str) -> None:
             inout_names.append(k)
             onnx_outputs.append(onnx_value(v, k))
             if idx < len(self.outputs):
+                if isinstance(self.outputs[idx], tuple):
+                    raise RuntimeError('Models returning nested lists/tuples are not supported yet')
                 _apply_tensor_info_to_value_info(onnx_outputs[-1], self.outputs[idx])
                 apply_dynamic_axes_info(onnx_outputs[-1], k)
 
@@ -878,11 +880,11 @@ def _convert(self) -> None:
             assert not to_utils.is_in_onnx_export()  # type: ignore[no-untyped-call]
             with to_utils.select_model_mode_for_export(self.original_model, self.training):
                 to_utils.__IN_ONNX_EXPORT = True
-                prev_opset_version = sym_hel._export_onnx_opset_version
+                prev_opset_version = GLOBALS.export_onnx_opset_version
                 sym_hel._set_opset_version(self.opset_version)  # type: ignore[no-untyped-call]
-                prev_export_type = sym_hel._operator_export_type
+                prev_export_type = GLOBALS.operator_export_type
                 sym_hel._set_operator_export_type(self.operator_export_type)  # type: ignore[no-untyped-call]
-                prev_shape_inference = sym_hel._onnx_shape_inference
+                prev_shape_inference = GLOBALS.onnx_shape_inference
                 sym_hel._set_onnx_shape_inference(  # type: ignore[no-untyped-call]
                     False  # TODO(twata): Use `self.onnx_shape_inference`
                 )

tests/pytorch_pfn_extras_tests/onnx_tests/test_export_testcase.py

@@ -11,12 +11,12 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-from torch.onnx.symbolic_helper import _default_onnx_opset_version
 
 from pytorch_pfn_extras.onnx import export
 from pytorch_pfn_extras.onnx import export_testcase
 from pytorch_pfn_extras.onnx import is_large_tensor
 from pytorch_pfn_extras.onnx import LARGE_TENSOR_DATA_THRESHOLD
+import pytorch_pfn_extras.onnx._constants
 from pytorch_pfn_extras.onnx.strip_large_tensor import _strip_large_tensor_tool_impl
 from pytorch_pfn_extras.onnx.unstrip_tensor import unstrip
 
@@ -45,7 +45,7 @@ def forward(self, x):
 
 def _get_output_dir(d, **kwargs):
     output_dir_base = 'out'
-    opset_ver = kwargs.get('opset_version', _default_onnx_opset_version)
+    opset_ver = kwargs.get('opset_version', pytorch_pfn_extras.onnx._constants.onnx_default_opset)
 
     output_dir = os.path.join(
         output_dir_base, 'opset{}'.format(opset_ver), d)