Merge pull request #806 from take-cheeze/pt2.2

Support torch 2.2

Author: linshokaku

.flexci/config.pbtxt

@@ -110,6 +110,22 @@ configs {
   }
 }
 
+configs {
+  key: "pytorch-pfn-extras.torch202-linux"
+  value {
+    requirement {
+      cpu: 8
+      gpu: 2
+      memory: 52
+      disk: 10
+    }
+    time_limit: {
+      seconds: 1800
+    }
+    command: "bash .flexci/linux/main-flexci.sh torch202"
+  }
+}
+
 configs {
   key: "pytorch-pfn-extras.torch110-win"
   value {
@@ -211,3 +227,20 @@ configs {
     command: ".flexci\\windows\\run.bat torch201"
   }
 }
+
+configs {
+  key: "pytorch-pfn-extras.torch202-win"
+  value {
+    requirement {
+      cpu: 4
+      gpu: 2
+      memory: 24
+      disk: 10
+      image: "windows"
+    }
+    time_limit: {
+      seconds: 1200
+    }
+    command: ".flexci\\windows\\run.bat torch202"
+  }
+}

.flexci/linux/build_and_push.sh

@@ -81,6 +81,15 @@ case "${TARGET}" in
             --build-arg pip_install_dep_args="cupy-cuda11x"
         ;;
 
+    torch202 )
+        # PyTorch 2.2 + Python 3.10
+        docker_build_and_push \
+            --build-arg base_image="nvidia/cuda:12.1.0-cudnn8-devel-ubuntu20.04" \
+            --build-arg python_version="3.10.5" \
+            --build-arg pip_install_torch_args="torch==2.2.* torchvision==0.17.* -f https://download.pytorch.org/whl/cu121/torch_stable.html" \
+            --build-arg pip_install_dep_args="cupy-cuda12x"
+        ;;
+
     * )
         echo "${1}: Unknown test name."
         exit 1

.flexci/windows/_flexci.ps1

@@ -52,6 +52,14 @@ function ActivateCUDA($version) {
         $Env:CUDA_PATH = $Env:CUDA_PATH_V11_8
     } elseif ($version -eq "11.x") {
         $Env:CUDA_PATH = $Env:CUDA_PATH_V11_8
+    } elseif ($version -eq "12.0") {
+        $Env:CUDA_PATH = $Env:CUDA_PATH_V12_0
+    } elseif ($version -eq "12.1") {
+        $Env:CUDA_PATH = $Env:CUDA_PATH_V12_1
+    } elseif ($version -eq "12.2") {
+        $Env:CUDA_PATH = $Env:CUDA_PATH_V12_2
+    } elseif ($version -eq "12.x") {
+        $Env:CUDA_PATH = $Env:CUDA_PATH_V12_2
     } else {
         throw "Unsupported CUDA version: $version"
     }

.flexci/windows/test.ps1

@@ -51,6 +51,13 @@ if ($test -eq "torch110") {
     RunOrDie python -m pip install -U pip "setuptools<59.6"
     RunOrDieWithRetry 3 python -m pip install torch==2.1.* torchvision==0.16.* -f https://download.pytorch.org/whl/cu118/torch_stable.html
 
+} elseif ($test -eq "torch202") {
+    # PyTorch 2.2 + Python 3.10
+    ActivateCUDA 12.1
+    ActivatePython 3.10
+    RunOrDie python -m pip install -U pip "setuptools<59.6"
+    RunOrDieWithRetry 3 python -m pip install torch==2.2.* torchvision==0.17.* -f https://download.pytorch.org/whl/cu121/torch_stable.html
+
 } else {
     throw "Unsupported test variant: $test"
 }

.github/workflows/pretest-and-test.yml

@@ -7,7 +7,7 @@ jobs:
     runs-on: ubuntu-22.04
     strategy:
       matrix:
-        torch: ['1.10.*', '1.11.*', '1.12.*', '1.13.*', '2.0.*', '2.1.*']
+        torch: ['1.10.*', '1.11.*', '1.12.*', '1.13.*', '2.0.*', '2.1.*', '2.2.*']
 
     steps:
     - name: Checkout

pytorch_pfn_extras/nn/modules/lazy_batchnorm.py

@@ -47,6 +47,7 @@ def forward(self, input: torch.Tensor) -> torch.Tensor:
                 )
             if self.track_running_stats:
                 assert isinstance(self.running_mean, torch.Tensor)
+                assert isinstance(self.num_features, int)
                 self.running_mean = torch.zeros(
                     self.num_features,
                     device=self.running_mean.device,

pytorch_pfn_extras/onnx/pfto_exporter/export.py

@@ -507,7 +507,7 @@ def optimize_onnx(self, graph: torch._C.Graph) -> torch._C.Graph:
             for k, t in folded.items():
                 c: torch._C.Value = graph.create("onnx::Constant", 1).output()
                 assert isinstance(t, torch.Tensor)
-                c.node().t_("value", cast(torch.Tensor, t))
+                c.node().t_("value", t)
                 graph.prependNode(c.node())
                 # TODO(twata): Determine folded nodes from original graph and document it
                 self.node_doc_string[c.node()] = f"Constant folded node: {input_table[k]}"
@@ -1155,9 +1155,9 @@ def _convert(self) -> None:
                     GLOBALS.onnx_shape_inference = False
                 else:
                     to_utils.__IN_ONNX_EXPORT = True  # type: ignore[attr-defined]
-                    sym_hel._set_opset_version(self.opset_version)  # type: ignore[no-untyped-call]
-                    sym_hel._set_operator_export_type(self.operator_export_type)  # type: ignore[no-untyped-call]
-                    sym_hel._set_onnx_shape_inference(  # type: ignore[no-untyped-call]
+                    sym_hel._set_opset_version(self.opset_version)  # type: ignore[attr-defined, no-untyped-call]
+                    sym_hel._set_operator_export_type(self.operator_export_type)  # type: ignore[attr-defined, no-untyped-call]
+                    sym_hel._set_onnx_shape_inference(  # type: ignore[attr-defined, no-untyped-call]
                         False  # TODO(twata): Use `self.onnx_shape_inference`
                     )
                 with record("pfto.original_outputs"):
@@ -1177,11 +1177,11 @@ def _convert(self) -> None:
             else:
                 to_utils.__IN_ONNX_EXPORT = False  # type: ignore[attr-defined]
                 if prev_opset_version is not None:
-                    sym_hel._set_opset_version(prev_opset_version)  # type: ignore[no-untyped-call]
+                    sym_hel._set_opset_version(prev_opset_version)  # type: ignore[attr-defined, no-untyped-call]
                 if prev_export_type is not None:
-                    sym_hel._set_operator_export_type(prev_export_type)  # type: ignore[no-untyped-call]
+                    sym_hel._set_operator_export_type(prev_export_type)  # type: ignore[attr-defined, no-untyped-call]
                 if prev_shape_inference is not None:
-                    sym_hel._set_onnx_shape_inference(prev_shape_inference)  # type: ignore[no-untyped-call]
+                    sym_hel._set_onnx_shape_inference(prev_shape_inference)  # type: ignore[attr-defined, no-untyped-call]
 
     def generate(self, f: Union[str, typing.IO]) -> None:
         with record("pfto.write_to_file"):

pytorch_pfn_extras/runtime/_to.py

@@ -58,10 +58,10 @@ def to(
     runtime = runtime_class(device, options)
     obj = module_or_tensor
     if isinstance(obj, torch.nn.Module):
-        obj = runtime.move_module(obj)
-        for module in obj.modules():
+        mod = runtime.move_module(obj)
+        for module in mod.modules():
             ppe.runtime._runtime._set_module_runtime_tag(module, runtime)
-        return obj
+        return mod
     elif isinstance(obj, torch.Tensor):
         return runtime.move_tensor(obj)
     else: