pytorch中的add_module函數

現只講在自定義網絡中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對模塊進行封裝等等。

class NeuralNetwork(nn.Module):
    def __init__(self):
        super(NeuralNetwork, self).__init__()
        self.layers = nn.Linear(28*28,28*28)
#         self.add_module('layers',nn.Linear(28*28,28*28))  #  跟上面的方式等價
        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

我們實例化類，然后輸出網絡的模塊看一下：

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

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

 

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

 

class NeuralNetwork(nn.Module):
    def __init__(self, layer_num):
        super(NeuralNetwork, self).__init__()
 4         self.layers = [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

對此我們再初始化一個實例，然后看下網絡中的模塊：

model = NeuralNetwork(2)
for index,item in enumerate(model.children()):
    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類型，所以不會自動注冊，那么就需要我們再定義后，手動注冊（下圖黃色標注部分）：

class NeuralNetwork(nn.Module):
    def __init__(self, layer_num):
        super(NeuralNetwork, self).__init__()
        self.layers = [nn.Linear(28*28,28*28) for _ in range(layer_num)]
        for i,layer in enumerate(self.layers):  6             self.add_module('layer_{}'.format(i),layer)  7         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

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

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