Repository: pfnet/pytorch-pfn-extras

Issue:
Feature request: LazyBatchNorm
In my project, I implemented LazyBatchNorm like this and I think pytorch-pfn-extras should have the similar feature:

```
class LazyBatchNorm1d(
    ppe.nn.modules.lazy.LazyInitializationMixin, torch.nn.BatchNorm1d
):
    """
    LazyBatchNorm1d is a lazy version of BatchNorm1d.  It does not require to
    specify the number of the input parameters beforehand.
    """

    lazy_parameter_names = ("weight", "bias")

    def __init__(self, num_features, *args, **kwargs):
        super().__init__(num_features or 0, *args, **kwargs)
        if num_features is None:
            self.num_features = None
            self.weight = ppe.nn.modules.lazy.UninitializedParameter()
            self.bias = ppe.nn.modules.lazy.UninitializedParameter()

    def forward(self, input):
        if isinstance(self.weight, ppe.nn.modules.lazy.UninitializedParameter):
            self.num_features = input.shape[-1]
            if self.affine:
                self.weight = torch.nn.Parameter(
                    self.weight.new_empty(self.num_features)
                )
                self.bias = torch.nn.Parameter(self.weight.new_empty(self.num_features))
            if self.track_running_stats:
                self.running_mean = torch.zeros(
                    self.num_features, device=self.running_mean.device
                )
                self.running_var = torch.ones(
                    self.num_features, device=self.running_var.device
                )
            self.reset_parameters()
        return super().forward(input)

    def reset_parameters(self):
        if self.lazy_parmeters_determined:
            super().reset_parameters()
```

Another discussion:
It would be good to have LazyBatchNorm instead of having LazyBatchNorm*d-s since lazy ones know its desired shapes. Any thoughts?


Discussion:
Thank you for your contirbution!

> It would be good to have LazyBatchNorm instead of having LazyBatchNorm*d-s since lazy ones know its desired shapes. Any thoughts?

I feel explicitness is one of the PyTorch's philosophy because they provide `BatchNormXd` (X=1,2,3) even though everything is implemented in hidden `_BatchNorm` class. https://pytorch.org/docs/master/_modules/torch/nn/modules/batchnorm.html
I think it's better to do as the Romans do.

Separate `LazyBatchNormXd` implementation will also make users easier to migrate from PPE's lazy implementation to PyTorch's upstream lazy implementation, which is under proposal in https://github.com/pytorch/pytorch/pull/40493
I will take care of this

Resolving patch:
diff --git a/pytorch_pfn_extras/nn/__init__.py b/pytorch_pfn_extras/nn/__init__.py
--- a/pytorch_pfn_extras/nn/__init__.py
+++ b/pytorch_pfn_extras/nn/__init__.py
@@ -2,4 +2,7 @@
 from pytorch_pfn_extras.nn.modules.lazy_conv import LazyConv1d  # NOQA
 from pytorch_pfn_extras.nn.modules.lazy_conv import LazyConv2d  # NOQA
 from pytorch_pfn_extras.nn.modules.lazy_conv import LazyConv3d  # NOQA
+from pytorch_pfn_extras.nn.modules.lazy_batchnorm import LazyBatchNorm1d  # NOQA
+from pytorch_pfn_extras.nn.modules.lazy_batchnorm import LazyBatchNorm2d  # NOQA
+from pytorch_pfn_extras.nn.modules.lazy_batchnorm import LazyBatchNorm3d  # NOQA
 from pytorch_pfn_extras.nn.modules.extended_sequential import ExtendedSequential  # NOQA
diff --git a/pytorch_pfn_extras/nn/modules/lazy_batchnorm.py b/pytorch_pfn_extras/nn/modules/lazy_batchnorm.py
new file mode 100644
--- /dev/null
+++ b/pytorch_pfn_extras/nn/modules/lazy_batchnorm.py
@@ -0,0 +1,79 @@
+import torch
+
+from pytorch_pfn_extras.nn.modules.lazy import LazyInitializationMixin
+from pytorch_pfn_extras.nn.modules.lazy import UninitializedParameter
+
+
+class _LazyBatchNorm(
+    LazyInitializationMixin,
+    torch.nn.modules.batchnorm._BatchNorm
+):
+
+    lazy_parameter_names = ('weight', 'bias')
+
+    def __init__(self, num_features, *args, **kwargs):
+        super().__init__(num_features or 0, *args, **kwargs)
+        if not self.affine:
+            raise ValueError(
+                'LazyBatchNorm is not compatible with affine=False.'
+                ' Use the regular BatchNorm layers instead')
+        # weight and bias are registered in the mixin
+        if num_features is None:
+            self.num_features = None
+            if self.track_running_stats:
+                # these buffers are not always needed
+                # so we avoid explicit initializations
+                self.lazy_buffer_names = ('running_mean', 'running_var')
+
+    def reset_parameters(self) -> None:
+        if self.lazy_parmeters_determined:
+            super().reset_parameters()
+
+    def forward(self, input):
+        if isinstance(self.weight, UninitializedParameter):
+            self.num_features = input.shape[-1]
+            if self.affine:
+                self.weight = torch.nn.Parameter(
+                    self.weight.new_empty(self.num_features)
+                )
+                self.bias = torch.nn.Parameter(
+                    self.weight.new_empty(self.num_features)
+                )
+            if self.track_running_stats:
+                self.running_mean = torch.zeros(
+                    self.num_features, device=self.running_mean.device,
+                    dtype=self.running_mean.dtype
+                )
+                self.running_var = torch.ones(
+                    self.num_features, device=self.running_var.device,
+                    dtype=self.running_mean.dtype
+                )
+            self.reset_parameters()
+        return super().forward(input)
+
+
+class LazyBatchNorm1d(_LazyBatchNorm, torch.nn.BatchNorm1d):
+    """BatchNorm1d module with lazy weight initialization.
+
+    When ``num_features`` is ``None``, it is determined at the first time of
+    the forward step.
+    """
+    pass
+
+
+class LazyBatchNorm2d(_LazyBatchNorm, torch.nn.BatchNorm2d):
+    """BatchNorm2d module with lazy weight initialization.
+
+    When ``num_features`` is ``None``, it is determined at the first time of
+    the forward step.
+    """
+    pass
+
+
+class LazyBatchNorm3d(_LazyBatchNorm, torch.nn.BatchNorm3d):
+    """BatchNorm3d module with lazy weight initialization.
+
+    When ``num_features`` is ``None``, it is determined at the first time of
+    the forward step.
+    """
+    pass


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_lazy_batchnorm.py b/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_lazy_batchnorm.py
new file mode 100644
--- /dev/null
+++ b/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_lazy_batchnorm.py
@@ -0,0 +1,43 @@
+import torch
+from torch import nn
+
+from pytorch_pfn_extras.nn import LazyBatchNorm1d, LazyBatchNorm2d, LazyBatchNorm3d  # NOQA
+
+from tests.pytorch_pfn_extras_tests.nn_tests.modules_tests.test_lazy import \
+    LazyTestBase
+
+
+class TestLazyBatchNorm1d(LazyTestBase):
+
+    def get_original_module(self):
+        return nn.BatchNorm1d(10)
+
+    def get_lazy_module(self):
+        return LazyBatchNorm1d(None)
+
+    def get_input(self):
+        return torch.rand(10, 10)
+
+
+class TestLazyBatchNorm2d(LazyTestBase):
+
+    def get_original_module(self):
+        return nn.BatchNorm2d(10)
+
+    def get_lazy_module(self):
+        return LazyBatchNorm2d(None)
+
+    def get_input(self):
+        return torch.rand(10, 10, 10, 10)
+
+
+class TestLazyBatchNorm3d(LazyTestBase):
+
+    def get_original_module(self):
+        return nn.BatchNorm3d(10)
+
+    def get_lazy_module(self):
+        return LazyBatchNorm3d(None)
+
+    def get_input(self):
+        return torch.rand(10, 10, 10, 10, 10)