Repository: pytorch/pytorch

Issue:
Add argument to children() to yield names of the modules as well
Currently, in order to selectively iterate through the modules in a network, we need to do something like
```python
for name, module in model._modules.items():
    if name in KEEP_LIST:
        ...
```
It would be good to add an option to `.children()` (and eventually to `.modules()` as well) to yield the `name` of the modules as well. Maybe something like `.children(named=True)`?


Discussion:

Resolving patch:
diff --git a/torch/nn/modules/module.py b/torch/nn/modules/module.py
--- a/torch/nn/modules/module.py
+++ b/torch/nn/modules/module.py
@@ -340,21 +340,72 @@ def parameters(self, memo=None):
                 yield p
 
     def children(self):
-        """Returns an iterator over children modules."""
+        """Returns an iterator over immediate children modules."""
+        for name, module in self.named_children():
+            yield module
+
+    def named_children(self):
+        """Returns an iterator over immediate children modules, yielding both
+        the name of the module as well as the module itself.
+
+        Example:
+            >>> for name, module in model.named_children():
+            >>>     if name in ['conv4', 'conv5']:
+            >>>         print(module)
+        """
         memo = set()
-        for module in self._modules.values():
+        for name, module in self._modules.items():
             if module is not None and module not in memo:
                 memo.add(module)
-                yield module
+                yield name, module
+
+    def modules(self):
+        """Returns an iterator over all modules in the network.
+
+        Note:
+            Duplicate modules are returned only once. In the following
+            example, ``l`` will be returned only once.
+
+            >>> l = nn.Linear(2, 2)
+            >>> net = nn.Sequential(l, l)
+            >>> for idx, m in enumerate(net.modules()):
+            >>>     print(idx, '->', m)
+            0 -> Sequential (
+              (0): Linear (2 -> 2)
+              (1): Linear (2 -> 2)
+            )
+            1 -> Linear (2 -> 2)
+        """
+        for name, module in self.named_modules():
+            yield module
+
+    def named_modules(self, memo=None, prefix=''):
+        """Returns an iterator over all modules in the network, yielding
+        both the name of the module as well as the module itself.
+
+        Note:
+            Duplicate modules are returned only once. In the following
+            example, ``l`` will be returned only once.
+
+            >>> l = nn.Linear(2, 2)
+            >>> net = nn.Sequential(l, l)
+            >>> for idx, m in enumerate(net.named_modules()):
+            >>>     print(idx, '->', m)
+            0 -> ('', Sequential (
+              (0): Linear (2 -> 2)
+              (1): Linear (2 -> 2)
+            ))
+            1 -> ('0', Linear (2 -> 2))
+        """
 
-    def modules(self, memo=None):
         if memo is None:
             memo = set()
         if self not in memo:
             memo.add(self)
-            yield self
-            for module in self.children():
-                for m in module.modules(memo):
+            yield prefix, self
+            for name, module in self._modules.items():
+                submodule_prefix = prefix + ('.' if prefix else '') + name
+                for m in module.named_modules(memo, submodule_prefix):
                     yield m
 
     def train(self, mode=True):


Test patch:
diff --git a/test/test_nn.py b/test/test_nn.py
--- a/test/test_nn.py
+++ b/test/test_nn.py
@@ -443,6 +443,29 @@ def __init__(self):
         self.assertEqual(num_params(n), 3)
         self.assertEqual(num_params(s), 3)
 
+    def test_children(self):
+        l1 = nn.Linear(2, 2)
+        l2 = nn.Linear(2, 2)
+        l3 = nn.Linear(2, 2)
+        l4 = nn.Linear(2, 2)
+        subnet = nn.Sequential(l3, l4)
+        s = nn.Sequential(l1, l2, l1, l2, subnet)
+        self.assertEqual(list(s.children()), [l1, l2, subnet])
+
+    def test_named_children(self):
+        l1 = nn.Linear(2, 2)
+        l2 = nn.Linear(2, 2)
+        l3 = nn.Linear(2, 2)
+        l4 = nn.Linear(2, 2)
+        subnet = nn.Sequential(l3, l4)
+        s = nn.Sequential()
+        s.add_module('layer1', l1)
+        s.add_module('layer2', l2)
+        s.add_module('layer3', l1)
+        s.add_module('layer4', l2)
+        s.add_module('subnet', subnet)
+        self.assertEqual(list(s.named_children()), [('layer1', l1), ('layer2', l2), ('subnet', subnet)])
+
     def test_modules(self):
         class Net(nn.Module):
             def __init__(self):
@@ -456,6 +479,26 @@ def __init__(self):
         s = nn.Sequential(n, n, n, n)
         self.assertEqual(list(s.modules()), [s, n, l])
 
+    def test_named_modules(self):
+        class Net(nn.Module):
+            def __init__(self):
+                super(Net, self).__init__()
+                self.l1 = l
+                self.l2 = l
+                self.param = Variable(torch.Tensor(3, 5))
+                self.block = block
+        l = nn.Linear(10, 20)
+        l1 = nn.Linear(10, 20)
+        l2 = nn.Linear(10, 20)
+        block = nn.Sequential()
+        block.add_module('linear1', l1)
+        block.add_module('linear2', l2)
+        n = Net()
+        s = nn.Sequential(n, n, n, n)
+        self.assertEqual(list(s.named_modules()), [('', s), ('0', n), ('0.l1', l),
+                                                   ('0.block', block), ('0.block.linear1', l1),
+                                                   ('0.block.linear2', l2)])
+
     def test_Sequential_getitem(self):
         l1 = nn.Linear(10, 20)
         l2 = nn.Linear(20, 30)