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對模塊進行封裝等等，

 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

標籤：其他

上一篇：升級openssl并重新編譯Nginx

下一篇：復雜多邊形的三角剖分