Merge branch 'master' into fix-test-torch-112

Author: asi1024

pytorch_pfn_extras/nn/modules/lazy.py

@@ -129,7 +129,7 @@ def _lazy_load_hook(  # type: ignore[no-untyped-def]
 
 class UninitializedParameter(torch.nn.Parameter):
 
-    def __repr__(self) -> str:
+    def __repr__(self) -> str:  # type: ignore[override]
         return 'Uninitialized lazy parameter'
 
     def share_memory_(self) -> 'UninitializedParameter':

pytorch_pfn_extras/nn/parallel/distributed.py

@@ -233,10 +233,10 @@ def forward(self, *args: Any, **kwargs: Any) -> Any:
 
     def load_state_dict(
             self,
-            state_dict: 'OrderedDict[str, torch.Tensor]',
+            state_dict: 'Mapping[str, torch.Tensor]',
             strict: bool = True,
     ) -> None:
-        self.module.load_state_dict(state_dict, strict=strict)
+        self.module.load_state_dict(state_dict, strict=strict)  # type: ignore[arg-type]
 
     T_destination = TypeVar('T_destination', bound=Mapping[str, torch.Tensor])
 

pytorch_pfn_extras/onnx/_constants.py

@@ -5,7 +5,7 @@
 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_default_opset = torch.onnx.symbolic_helper._default_onnx_opset_version  # type: ignore
+    onnx_main_opset = torch.onnx.symbolic_helper._onnx_main_opset  # type: ignore
+    onnx_stable_opsets = torch.onnx.symbolic_helper._onnx_stable_opsets  # type: ignore
     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

@@ -1,29 +1,30 @@
 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
+if pytorch_pfn_extras.requires("1.12.0"):
+    import torch.onnx._globals
+    GLOBALS = torch.onnx._globals.GLOBALS
 
-    @property
-    def operator_export_type(self) -> Optional[torch._C._onnx.OperatorExportTypes]:
-        return torch.onnx.symbolic_helper._operator_export_type
+else:
+    import torch.onnx.symbolic_helper as symhel
 
-    @property
-    def training_mode(self) -> Optional[torch._C._onnx.TrainingMode]:
-        return torch.onnx.symbolic_helper._training_mode
+    class _InternalGlobalsBeforeTorch1_11:
+        @property
+        def export_onnx_opset_version(self) -> int:
+            return symhel._export_onnx_opset_version  # type: ignore
 
-    @property
-    def onnx_shape_inference(self) -> bool:
-        return torch.onnx.symbolic_helper._onnx_shape_inference
+        @property
+        def operator_export_type(self) -> Optional[torch._C._onnx.OperatorExportTypes]:
+            return symhel._operator_export_type  # type: ignore
 
+        @property
+        def training_mode(self) -> Optional[torch._C._onnx.TrainingMode]:
+            return symhel._training_mode  # type: ignore
 
-if pytorch_pfn_extras.requires("1.12.0"):
-    import torch.onnx._globals
-    GLOBALS = torch.onnx._globals.GLOBALS
-else:
-    GLOBALS = _InternalGlobalsBeforeTorch1_11()
+        @property
+        def onnx_shape_inference(self) -> bool:
+            return symhel._onnx_shape_inference  # type: ignore
+
+    GLOBALS = _InternalGlobalsBeforeTorch1_11()  # type: ignore[assignment]

pytorch_pfn_extras/onnx/export_testcase.py

@@ -111,7 +111,12 @@ def _export_util(
         operator_export_type = OperatorExportTypes.ONNX_ATEN if\
             aten else OperatorExportTypes.RAW  # type: ignore
     elif operator_export_type is None:
-        if torch.onnx.PYTORCH_ONNX_CAFFE2_BUNDLE:
+        if pytorch_pfn_extras.requires("1.12.0"):
+            use_onnx_aten_fallback = torch.onnx._CAFFE2_ATEN_FALLBACK  # type: ignore[attr-defined]
+        else:
+            use_onnx_aten_fallback = torch.onnx.PYTORCH_ONNX_CAFFE2_BUNDLE  # type: ignore[attr-defined]
+
+        if use_onnx_aten_fallback:
             operator_export_type = OperatorExportTypes.ONNX_ATEN_FALLBACK
         else:
             operator_export_type = OperatorExportTypes.ONNX

pytorch_pfn_extras/utils/comparer.py

@@ -12,6 +12,7 @@
 import torch.nn
 import torch.testing
 
+import pytorch_pfn_extras
 from pytorch_pfn_extras import handler as _handler_module
 from pytorch_pfn_extras.handler import _logic
 from pytorch_pfn_extras.training import _trainer
@@ -191,8 +192,13 @@ def compare_fn(
             val1 = val1.cpu().detach()
         if isinstance(val2, torch.Tensor):
             val2 = val2.cpu().detach()
-        torch.testing.assert_allclose(
-            val1, val2, rtol=rtol, atol=atol, equal_nan=equal_nan)
+
+        if pytorch_pfn_extras.requires("1.9.0"):
+            assert_close = torch.testing.assert_close  # type: ignore[attr-defined]
+        else:
+            assert_close = torch.testing.assert_allclose  # type: ignore[assignment]
+
+        assert_close(val1, val2, rtol=rtol, atol=atol, equal_nan=equal_nan)
 
     return compare_fn
 

tests/pytorch_pfn_extras_tests/test_handler.py

@@ -6,6 +6,13 @@
 import pytorch_pfn_extras as ppe
 
 
+def torch_testing_assert_close(*args, **kwargs):
+    if ppe.requires("1.10.0"):
+        torch.testing.assert_close(*args, **kwargs)
+    else:
+        torch.testing.assert_allclose(*args, **kwargs)
+
+
 class MockRuntime(ppe.runtime.BaseRuntime):
     def __init__(self, device, options):
         super().__init__(device, options)
@@ -361,8 +368,8 @@ def test_train_step(self):
         model = models['main']
         assert input.grad is not None
         # The gradient of a linear layer is its transposed weight
-        torch.testing.assert_allclose(input.grad, model.weight.T)
-        torch.testing.assert_allclose(out, model(input))
+        torch_testing_assert_close(input.grad, model.weight.T)
+        torch_testing_assert_close(out, model(input))
 
     @pytest.mark.parametrize(
         'to_backprop',
@@ -403,11 +410,11 @@ def forward(self, x):
         grad = torch.zeros(1)
         for val in to_backprop:
             grad = grad + getattr(model, f'l{val}').weight.T
-        torch.testing.assert_allclose(input.grad, grad)
+        torch_testing_assert_close(input.grad, grad)
 
         # Check that logic step does not change the value of weight
         for val in original_parameters:
-            torch.testing.assert_allclose(
+            torch_testing_assert_close(
                 original_parameters[val], getattr(model, f'l{val}').weight)
 
     def test_train_step_backward_nograd(self):
@@ -461,7 +468,7 @@ def test_train_step_optimizers(self):
         w_grad = model.weight.grad.clone().detach()
         logic.train_step_optimizers(model, optimizers, 0)
         # Checks that the value was correctly updated
-        torch.testing.assert_allclose(m_weight - w_grad, model.weight.T)
+        torch_testing_assert_close(m_weight - w_grad, model.weight.T)
 
     @pytest.mark.gpu
     def test_grad_scaler(self):
@@ -473,12 +480,12 @@ def test_grad_scaler(self):
         m_weight = model.weight.clone().detach()
         w_grad = model.weight.grad.clone().detach()
         # The gradient of a linear layer is its transposed weight
-        torch.testing.assert_allclose(input.grad, scaler.scale(model.weight.T))
-        torch.testing.assert_allclose(out, model(input))
+        torch_testing_assert_close(input.grad, scaler.scale(model.weight.T))
+        torch_testing_assert_close(out, model(input))
         logic.train_step_optimizers(model, optimizers, 0)
         # Checks that the value was correctly updated and gradients deescaled
         # before the update
-        torch.testing.assert_allclose(
+        torch_testing_assert_close(
             scaler.scale(m_weight) - w_grad, scaler.scale(model.weight.T))
 
     @pytest.mark.gpu
@@ -513,4 +520,4 @@ def test_eval_step(self):
         models = {'main': model}
         models['main'].eval()
         out = logic.eval_step(models, 0, input)
-        torch.testing.assert_allclose(out, model(input))
+        torch_testing_assert_close(out, model(input))