現只講在自定義網路中add_module的作用。

總結：

在自定義網路的時候，由於自定義變數不是Module型別（例如，我們用List封裝了幾個網路），所以pytorch不會自動註冊網路模組。add_module函式用來為網路新增模組的，所以我們可以使用這個函式手動新增自定義的網路模組。當然，這種情況，我們也可以使用ModuleList來封裝自定義模組，pytorch就會自動註冊了。

Let't start!

add_module函式是在自定義網路新增子模組，例如，當我們自定義一個網路膚過程中，我們既可以

（1）通過self.module=xxx_module的方式（如下面第3行程式碼），新增網路模組；

（2）通過add_module函式對網路中新增模組。

（3）通過用nn.Sequential對模組進行封裝等等。

 1 class NeuralNetwork(nn.Module):
 2     def __init__(self):
 3         super(NeuralNetwork, self).__init__()
 4         self.layers = nn.Linear(28*28,28*28)
 5 #         self.add_module('layers',nn.Linear(28*28,28*28))  #  跟上面的方式等價
 6         self.linear_relu_stack = nn.Sequential(
 
 7             nn.Linear(28*28, 512),
 8             nn.ReLU()
 9         )
10 
11     def forward(self, x):
12         for layer in layers:
13             x = layer(x)
14         logits = self.linear_relu_stack(x)
15         return logits

我們例項化類，然後輸出網路的模組看一下：

1 0 Linear(in_features=784, out_features=784, bias=True)
 
2 1 Sequential(
3   (0): Linear(in_features=784, out_features=512, bias=True)
4   (1): ReLU()
5 )

會發現，上面定義的網路子模組都有：Linear和Sequential。

但是，有時候pytorch不會自動給我們註冊模組，我們需要根據傳進來的引數對網路進行初始化，例如：

 1 class NeuralNetwork(nn.Module):
 2     def __init__(self, layer_num):
 3         super(NeuralNetwork, self).__init__()
 4         self.layers = [nn.Linear(28*28,28*28) for _ in range(layer_num)]
 5         self.linear_relu_stack = nn.Sequential(
 6             nn.Linear(28*28, 512),
 7             nn.ReLU()
 8         )
 9 
10     def forward(self, x):
11         for layer in layers:
12             x = layer(x)
13         logits = self.linear_relu_stack(x)
14         return logits

對此我們再初始化一個例項，然後看下網路中的模組：

1 model = NeuralNetwork(2)
2 for index,item in enumerate(model.children()):
3     print(index,item)

輸出結果就是：

0 Sequential(
  (0): Linear(in_features=784, out_features=512, bias=True)
  (1): ReLU()
)　

你會發現定義的Linear模組都不見了，而上面定義的時候，明明都制訂了。這是因為pytorch在註冊模組的時候，會檢視成員的型別，如果成員變數型別是Module的子類，那麼pytorch就會註冊這個模組，否則就不會。

這裡的self.layers是python中的List型別，所以不會自動註冊，那麼就需要我們再定義後，手動註冊（下圖黃色標註部分）：

 1 class NeuralNetwork(nn.Module):
 2     def __init__(self, layer_num):
 3         super(NeuralNetwork, self).__init__()
 4         self.layers = [nn.Linear(28*28,28*28) for _ in range(layer_num)]
 5         for i,layer in enumerate(self.layers):
 6             self.add_module('layer_{}'.format(i),layer)
 7         self.linear_relu_stack = nn.Sequential(
 8             nn.Linear(28*28, 512),
 9             nn.ReLU()
10         )
11 
12     def forward(self, x):
13         for layer in layers:
14             x = layer(x)
15         logits = self.linear_relu_stack(x)
16         return logits

這樣我們再輸出模型的子模組的時候，就會得到：

model = NeuralNetwork(4)
for index,item in enumerate(model.children()):
    print(index,item)

# output
#0 Linear(in_features=784, out_features=784, bias=True)
#1 Linear(in_features=784, out_features=784, bias=True)
#2 Linear(in_features=784, out_features=784, bias=True)
#3 Linear(in_features=784, out_features=784, bias=True)
#4 Sequential(
#  (0): Linear(in_features=784, out_features=512, bias=True)
#  (1): ReLU()
#)

就會看到，已經有了自己註冊的模組。

當然，也可能覺得這種方式比較麻煩，每次都要自己註冊下，那能不能有一個類似List的類，在定義的時候就封裝一下呢？

可以，使用nn.ModuleList封裝一下即可達到相同的效果。

class NeuralNetwork(nn.Module):
    def __init__(self, layer_num):
        super(NeuralNetwork, self).__init__()
        self.layers = nn.ModuleList([nn.Linear(28*28,28*28) for _ in range(layer_num)])
        self.linear_relu_stack = nn.Sequential(
            nn.Linear(28*28, 512),
            nn.ReLU()
        )

    def forward(self, x):
        for layer in layers:
            x = layer(x)
        logits = self.linear_relu_stack(x)
        return logits

參考：
1. 部落格THE PYTORCH ADD_MODULE() FUNCTION link
2. pytorch 官方文件 中文連結English version