Merge pull request #23 from calvinmccarter-at-lightmatter/master

Fix _LazyBatchNorm import broken by torch 1.10.0

Author: Talmaj

onnx2pytorch/operations/batchnorm.py

@@ -1,7 +1,19 @@
 import warnings
 
 from torch import nn
-from torch.nn.modules.batchnorm import _BatchNorm, _LazyBatchNorm
+from torch.nn.modules.batchnorm import _BatchNorm
+
+try:
+    from torch.nn.modules.batchnorm import _LazyNormBase
+
+    class _LazyBatchNorm(_LazyNormBase, _BatchNorm):
+
+        cls_to_become = _BatchNorm
+
+
+except ImportError:
+    # for torch < 1.10.0
+    from torch.nn.modules.batchnorm import _LazyBatchNorm
 
 
 class LazyBatchNormUnsafe(_LazyBatchNorm):

onnx2pytorch/operations/instancenorm.py

@@ -2,71 +2,80 @@
 import torch
 
 from torch.nn.modules.instancenorm import _InstanceNorm
-from torch.nn.modules.lazy import LazyModuleMixin
-from torch.nn.parameter import UninitializedBuffer, UninitializedParameter
-
-
-class _LazyInstanceNorm(LazyModuleMixin, _InstanceNorm):
-
-    weight: UninitializedParameter  # type: ignore[assignment]
-    bias: UninitializedParameter  # type: ignore[assignment]
-
-    cls_to_become = _InstanceNorm
-
-    def __init__(
-        self,
-        eps=1e-5,
-        momentum=0.1,
-        affine=True,
-        track_running_stats=True,
-        device=None,
-        dtype=None,
-    ) -> None:
-        factory_kwargs = {"device": device, "dtype": dtype}
-        super(_LazyInstanceNorm, self).__init__(
-            # affine and track_running_stats are hardcoded to False to
-            # avoid creating tensors that will soon be overwritten.
-            0,
-            eps,
-            momentum,
-            False,
-            False,
-            **factory_kwargs,
-        )
-        self.affine = affine
-        self.track_running_stats = track_running_stats
-        if self.affine:
-            self.weight = UninitializedParameter(**factory_kwargs)
-            self.bias = UninitializedParameter(**factory_kwargs)
-        if self.track_running_stats:
-            self.running_mean = UninitializedBuffer(**factory_kwargs)
-            self.running_var = UninitializedBuffer(**factory_kwargs)
-            self.num_batches_tracked = torch.tensor(
-                0,
-                dtype=torch.long,
-                **{k: v for k, v in factory_kwargs.items() if k != "dtype"},
-            )
 
-    def reset_parameters(self) -> None:
-        if not self.has_uninitialized_params() and self.num_features != 0:
-            super().reset_parameters()
+try:
+    from torch.nn.modules.batchnorm import _LazyNormBase
+
+    class _LazyInstanceNorm(_LazyNormBase, _InstanceNorm):
+
+        cls_to_become = _InstanceNorm
+
+
+except ImportError:
+    from torch.nn.modules.lazy import LazyModuleMixin
+    from torch.nn.parameter import UninitializedBuffer, UninitializedParameter
 
-    def initialize_parameters(self, input) -> None:  # type: ignore[override]
-        if self.has_uninitialized_params():
-            self.num_features = input.shape[1]
+    class _LazyInstanceNorm(LazyModuleMixin, _InstanceNorm):
+
+        weight: UninitializedParameter  # type: ignore[assignment]
+        bias: UninitializedParameter  # type: ignore[assignment]
+
+        cls_to_become = _InstanceNorm
+
+        def __init__(
+            self,
+            eps=1e-5,
+            momentum=0.1,
+            affine=True,
+            track_running_stats=True,
+            device=None,
+            dtype=None,
+        ) -> None:
+            factory_kwargs = {"device": device, "dtype": dtype}
+            super(_LazyInstanceNorm, self).__init__(
+                # affine and track_running_stats are hardcoded to False to
+                # avoid creating tensors that will soon be overwritten.
+                0,
+                eps,
+                momentum,
+                False,
+                False,
+                **factory_kwargs,
+            )
+            self.affine = affine
+            self.track_running_stats = track_running_stats
             if self.affine:
-                assert isinstance(self.weight, UninitializedParameter)
-                assert isinstance(self.bias, UninitializedParameter)
-                self.weight.materialize((self.num_features,))
-                self.bias.materialize((self.num_features,))
+                self.weight = UninitializedParameter(**factory_kwargs)
+                self.bias = UninitializedParameter(**factory_kwargs)
             if self.track_running_stats:
-                self.running_mean.materialize(
-                    (self.num_features,)
-                )  # type:ignore[union-attr]
-                self.running_var.materialize(
-                    (self.num_features,)
-                )  # type:ignore[union-attr]
-            self.reset_parameters()
+                self.running_mean = UninitializedBuffer(**factory_kwargs)
+                self.running_var = UninitializedBuffer(**factory_kwargs)
+                self.num_batches_tracked = torch.tensor(
+                    0,
+                    dtype=torch.long,
+                    **{k: v for k, v in factory_kwargs.items() if k != "dtype"},
+                )
+
+        def reset_parameters(self) -> None:
+            if not self.has_uninitialized_params() and self.num_features != 0:
+                super().reset_parameters()
+
+        def initialize_parameters(self, input) -> None:  # type: ignore[override]
+            if self.has_uninitialized_params():
+                self.num_features = input.shape[1]
+                if self.affine:
+                    assert isinstance(self.weight, UninitializedParameter)
+                    assert isinstance(self.bias, UninitializedParameter)
+                    self.weight.materialize((self.num_features,))
+                    self.bias.materialize((self.num_features,))
+                if self.track_running_stats:
+                    self.running_mean.materialize(
+                        (self.num_features,)
+                    )  # type:ignore[union-attr]
+                    self.running_var.materialize(
+                        (self.num_features,)
+                    )  # type:ignore[union-attr]
+                self.reset_parameters()
 
 
 class LazyInstanceNormUnsafe(_LazyInstanceNorm):

tox.ini

@@ -4,14 +4,15 @@
 # and then run "tox" from this directory.
 
 [tox]
-envlist = clean,py36,py37,py38
+envlist = clean,py36,py37,py38,py38-torch19,py39
 
 [testenv]
 passenv =
     CIRCLE*
     KMP_DUPLICATE_LIB_OK
 deps =
     -rrequirements.txt
+	torch19: torch <= 1.9.0.
     pytest-cov
 commands =
     pytest --cov --cov-append --cov-report term --cov-report html tests/