【Pytorch】Module里的children()与modules()的区别

概述

children()与modules()都是返回网络模型里的组成元素，但是children()返回的是最外层的元素，modules()返回的是所有的元素，包括不同级别的子元素。

以代码为例

import torch
from torch import nn

# hyper parameters
in_dim=1
n_hidden_1=1
n_hidden_2=1
out_dim=1

class Net(nn.Module):
    def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
        super().__init__()

        self.layer = nn.Sequential(
            nn.Linear(in_dim, n_hidden_1), 
            nn.ReLU(True)
             )
        self.layer2 = nn.Sequential(
                nn.Linear(n_hidden_1, n_hidden_2),
                nn.ReLU(True),
            )
        self.layer3 = nn.Linear(n_hidden_2, out_dim)
        # print(self.modules())
        print("children")
        
        for i, module in enumerate( self.children()):
            print(i, module)
        print("modules")
        for i, module in enumerate( self.modules()):
            print(i, module)

    
    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        x = self.layer3(x)
        return x
      
        
model = Net(in_dim, n_hidden_1, n_hidden_2, out_dim)

这是一个三层的网络结构，将第一层的线性层和激活层放在一个nn.Sequential层中，将第二层的线性层和激活函数放在第二个nn.Sequential中，最后一个线性层作为单独第三层。

整个网络结构如下图所示：

接下来看一下代码__init__的print函数的打印信息：

• self.modules()

children
0 Sequential(
(0): Linear(in_features=1, out_features=1, bias=True)
(1): ReLU(inplace)
)
1 Sequential(
(0): Linear(in_features=1, out_features=1, bias=True)
(1): ReLU(inplace)
)
2 Linear(in_features=1, out_features=1, bias=True)

可以看出，self.children()存储网络结构的子层模块，也就是net's children那一层。

• self.modules()

modules
0 Net(
(layer): Sequential(
(0): Linear(in_features=1, out_features=1, bias=True)
(1): ReLU(inplace)
)
(layer2): Sequential(
(0): Linear(in_features=1, out_features=1, bias=True)
(1): ReLU(inplace)
)
(layer3): Linear(in_features=1, out_features=1, bias=True)
)
1 Sequential(
(0): Linear(in_features=1, out_features=1, bias=True)
(1): ReLU(inplace)
)
2 Linear(in_features=1, out_features=1, bias=True)
3 ReLU(inplace)
4 Sequential(
(0): Linear(in_features=1, out_features=1, bias=True)
(1): ReLU(inplace)
)
5 Linear(in_features=1, out_features=1, bias=True)
6 ReLU(inplace)
7 Linear(in_features=1, out_features=1, bias=True)

可以看出，self.modules()采用深度优先遍历的方式，存储了net的所有模块，包括net itself,net's children, children of net's children。

用list举例就是：

a=[1,2,[3,4]]
# children返回：
1,2，[3，4]
# modules返回：
[1,2,[3,4]], 1, 2, [3,4], 3, 4

参考博客

pytorch系列8 --self.modules() 和 self.children()的区别
pytorch Module里的children()与modules()的区别