Repository: pfnet/pytorch-pfn-extras

Issue:
Bug in sequential repeat when the layer has no parameters
**problems Statement**
In `init` mode, the repeated layer will be reset. In `Pytorch`, we used `reset_parameters` function to reset the parameters of layers, as here:
https://github.com/pfnet/pytorch-pfn-extras/blob/92dad97e3a3eaafb9eb2b2b804266b5ef922e4c7/pytorch_pfn_extras/nn/modules/extended_sequential.py#L15

However, there are layers have no parameters nor `reset_parameters`, such as `torch.nn.ReLU`. An error is raised when the model contains such layer.

**Error Message**
```
pytorch_pfn_extras/nn/modules/extended_sequential.py:68: in repeat
    model_list.append(self._copy_model(mode))
pytorch_pfn_extras/nn/modules/extended_sequential.py:27: in _copy_model
    return _reset_parameters(copy.deepcopy(self))
pytorch_pfn_extras/nn/modules/extended_sequential.py:9: in _reset_parameters
    _reset_parameters(submodel)
pytorch_pfn_extras/nn/modules/extended_sequential.py:17: in _reset_parameters
    model.reset_parameters()
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = ReLU(), name = 'reset_parameters'

    def __getattr__(self, name):
        if '_parameters' in self.__dict__:
            _parameters = self.__dict__['_parameters']
            if name in _parameters:
                return _parameters[name]
        if '_buffers' in self.__dict__:
            _buffers = self.__dict__['_buffers']
            if name in _buffers:
                return _buffers[name]
        if '_modules' in self.__dict__:
            modules = self.__dict__['_modules']
            if name in modules:
                return modules[name]
        raise AttributeError("'{}' object has no attribute '{}'".format(
>           type(self).__name__, name))
E       AttributeError: 'ReLU' object has no attribute 'reset_parameters'
```


Discussion:

Resolving patch:
diff --git a/pytorch_pfn_extras/nn/modules/extended_sequential.py b/pytorch_pfn_extras/nn/modules/extended_sequential.py
--- a/pytorch_pfn_extras/nn/modules/extended_sequential.py
+++ b/pytorch_pfn_extras/nn/modules/extended_sequential.py
@@ -1,5 +1,6 @@
 import torch
 import copy
+import warnings
 
 
 def _reset_parameters(model):
@@ -11,8 +12,19 @@ def _reset_parameters(model):
         for submodel in model.values():
             _reset_parameters(submodel)
     else:
-        if isinstance(model, torch.nn.Module):
+        if hasattr(model, 'reset_parameters'):
             model.reset_parameters()
+        elif hasattr(model, '_reset_parameters'):
+            model._reset_parameters()
+        else:
+            if (len(list(model.parameters())) != 0 or
+                    len(list(model.buffers())) != 0):
+                warnings.warn('Cannot reset the parameters of module {}. '
+                              'Consider adding `reset_parameters` or '
+                              '`_reset_parameters` '
+                              'functions to the module'.format(model),
+                              UserWarning)
+
     return model
 
 
@@ -39,6 +51,18 @@ def repeat(self, n_repeat: int, mode: 'str' = 'init'):
         The functions is supposed to behave the same way as `repeat`
         in `chainer`.
 
+        When the mode is set to ``init``, the default value,
+        modules will be copied and reinitialized by calling
+        ``reset_parameters`` (or ``_reset_parameters``) method.
+
+        To repeat user-defined modules, which have parameters or buffers,
+        with mode=``init`` in this Sequential,
+        you need to implement the ``reset_parameters`` or ``_reset_parameters``
+        method to the module to reinitialize parameters
+        and (if necessary) buffers;
+        otherwise the initialization cannot be performed
+        and a warning message will be shown.
+
         Args:
             n_repeat (int): Number of times to repeat.
             mode (str): It should be either ``init``, ``copy``, or ``share``.


Test patch:
diff --git a/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_extended_sequential.py b/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_extended_sequential.py
--- a/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_extended_sequential.py
+++ b/tests/pytorch_pfn_extras_tests/nn_tests/modules_tests/test_extended_sequential.py
@@ -1,38 +1,66 @@
 import unittest
 import pytest
+import functools
 
 import numpy
+import torch
 from torch import nn
 import pytorch_pfn_extras as ppe
 
 assertions = unittest.TestCase('__init__')
 
 
-@pytest.mark.parametrize('module', [nn.Sequential,
-                                    nn.ModuleList,
-                                    nn.ModuleDict])
+class UserDefinedLayer(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self):
+        pass
+
+
+@pytest.mark.parametrize('container', [
+    nn.Sequential,
+    nn.ModuleList,
+    nn.ModuleDict,
+])
+@pytest.mark.parametrize('irregular_layer', [
+    UserDefinedLayer,
+    # No reset_parameters
+    nn.ReLU,
+    # use reset_running_stats
+    functools.partial(
+        nn.BatchNorm1d, 1),
+    # use _reset_parameters
+    functools.partial(
+        nn.MultiheadAttention, 1, 1),
+    # ppe.nn layer
+    functools.partial(
+        ppe.nn.LazyConv1d, None, 1, 1),
+])
 class TestExtendedSequential(object):
 
     @pytest.fixture(autouse=True)
-    def setUp(self, module):
+    def setUp(self, container, irregular_layer):
         self.l1 = ppe.nn.LazyLinear(None, 3)
         self.l2 = nn.Linear(3, 2)
         self.l3 = nn.Linear(2, 3)
+        # a layer without reset_parameters
+        self.l4 = irregular_layer()
         # s1: l1 -> l2
-        if module == nn.Sequential:
-            self.s1 = module(self.l1, self.l2)
-        elif module == nn.ModuleDict:
-            self.s1 = module({
+        if container == nn.Sequential:
+            self.s1 = container(self.l1, self.l2)
+        elif container == nn.ModuleDict:
+            self.s1 = container({
                 'l1': self.l1,
                 'l2': self.l2})
         else:
-            self.s1 = module([self.l1, self.l2])
-        self.module = module
-        # s2: s1 (l1 -> l2) -> l3
-        self.s2 = ppe.nn.ExtendedSequential(self.s1, self.l3)
+            self.s1 = container([self.l1, self.l2])
+        self.container = container
+        # s2: s1 (l1 -> l2) -> l3 -> l4
+        self.s2 = ppe.nn.ExtendedSequential(self.s1, self.l3, self.l4)
 
     def test_repeat_with_init(self):
-        # s2 ((l1 -> l2) -> l3) -> s2 ((l1 -> l2) -> l3)
+        # s2 ((l1 -> l2) -> l3 -> l4) -> s2 ((l1 -> l2) -> l3 -> l4)
         ret = self.s2.repeat(2)
         assertions.assertIsNot(ret[0], self.s2)
         assertions.assertIs(type(ret[0]), type(self.s2))
@@ -40,7 +68,7 @@ def test_repeat_with_init(self):
         assertions.assertIs(type(ret[1]), type(self.s2))
 
         # bias is filled with 0, so they should have the same values
-        if self.module == nn.ModuleDict:
+        if self.container == nn.ModuleDict:
             numpy.testing.assert_array_equal(
                 ret[0][0]['l1'].bias.detach().numpy(),
                 ret[1][0]['l1'].bias.detach().numpy())
@@ -70,7 +98,7 @@ def test_repeat_with_init(self):
         assertions.assertEqual(len(ret), 0)
 
     def test_repeat_with_copy(self):
-        # s2 ((l1 -> l2) -> l3) -> s2 ((l1 -> l2) -> l3)
+        # s2 ((l1 -> l2) -> l3 -> l4) -> s2 ((l1 -> l2) -> l3 -> l4)
         ret = self.s2.repeat(2, mode='copy')
         assertions.assertIsNot(ret[0], self.s2)
         assertions.assertIs(type(ret[0]), type(self.s2))
@@ -79,7 +107,7 @@ def test_repeat_with_copy(self):
         assertions.assertIsNot(ret[0], ret[1])
 
         # b is filled with 0, so they should have the same values
-        if self.module == nn.ModuleDict:
+        if self.container == nn.ModuleDict:
             numpy.testing.assert_array_equal(
                 ret[0][0]["l1"].bias.detach().numpy(),
                 ret[1][0]["l1"].bias.detach().numpy())
@@ -106,7 +134,7 @@ def test_repeat_with_copy(self):
         assertions.assertEqual(len(ret), 0)
 
     def test_repeat_with_share(self):
-        # s2 ((l1 -> l2) -> l3) -> s2 ((l1 -> l2) -> l3)
+        # s2 ((l1 -> l2) -> l3 -> l4) -> s2 ((l1 -> l2) -> l3 -> l4)
         ret = self.s2.repeat(2, mode='share')
         assertions.assertIsNot(ret[0], self.s2)
         assertions.assertIs(type(ret[0]), type(self.s2))
@@ -114,7 +142,7 @@ def test_repeat_with_share(self):
         assertions.assertIs(type(ret[1]), type(self.s2))
 
         # b is filled with 0, so they should have the same values
-        if self.module == nn.ModuleDict:
+        if self.container == nn.ModuleDict:
             numpy.testing.assert_array_equal(
                 ret[0][0]["l1"].bias.detach().numpy(),
                 ret[1][0]["l1"].bias.detach().numpy())
@@ -136,3 +164,72 @@ def test_repeat_with_share(self):
         assertions.assertEqual(len(ret), 2)
         ret = self.s2.repeat(0, mode='share')
         assertions.assertEqual(len(ret), 0)
+
+
+class UserDefinedLayerWithReset(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self):
+        pass
+
+    def reset_parameters(self):
+        pass
+
+
+class UserDefinedLayerWithUnderScoreReset(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self):
+        pass
+
+    def _reset_parameters(self):
+        pass
+
+
+class UserDefinedLayerWithParameters(nn.Module):
+    def __init__(self):
+        super().__init__()
+        param = nn.Parameter(torch.zeros(1, 1))
+        self.register_parameter('weight', param)
+
+    def forward(self):
+        pass
+
+
+class UserDefinedLayerWithBuffer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer('weight', torch.zeros(1, 1))
+
+    def forward(self):
+        pass
+
+
+@pytest.mark.parametrize('module', [
+    # buit-in, no parameters
+    nn.ReLU,
+    # no parameters
+    UserDefinedLayer,
+    # has `_reset_parameters`
+    UserDefinedLayerWithUnderScoreReset,
+    # has `reset_parameters`
+    UserDefinedLayerWithReset,
+])
+def test_no_warning_when_repeat(module):
+    model = ppe.nn.ExtendedSequential(module())
+    # no warnings are raised on these modules
+    with pytest.warns(None):
+        model.repeat(2)
+
+
+@pytest.mark.parametrize('module', [
+    UserDefinedLayerWithParameters,
+    UserDefinedLayerWithBuffer,
+])
+def test_warning_when_repeat(module):
+    model = ppe.nn.ExtendedSequential(module())
+    # warnings are raised on these modules
+    with pytest.warns(UserWarning):
+        model.repeat(2)