pytorch提取神經網絡模型層結構和參數初始化

torch.nn.Module()類有一些重要屬性，我們可用其下面幾個屬性來實現對神經網絡層結構的提取：

torch.nn.Module.children()
torch.nn.Module.modules()
torch.nn.Module.named_children()
torch.nn.Module.named_moduless()

為方面說明，我們首先搭建一個簡單的神經網絡模型，后面所有的內容都是基於這個模型展開的。

import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()                                  # input shape (1, 28, 28)
        self.conv1 = nn.Sequential()
        self.conv1.add_module('conv1', nn.Conv2d(1, 16, 5, 1, 2))    # output shape (16, 28, 28)
        self.conv1.add_module('ReLU1', nn.ReLU())
        self.conv1.add_module('pool1', nn.MaxPool2d(2))              # output shape (16, 14, 14)

        self.conv2 = nn.Sequential()
        self.conv2.add_module('conv2', nn.Conv2d(16, 32, 5, 1, 2))   # output shape (32, 14, 14)
        self.conv2.add_module('ReLU2', nn.ReLU())
        self.conv2.add_module('pool2', nn.MaxPool2d(2))              # output shape (32, 7, 7)

        self.linear = nn.Linear(32 * 7 * 7, 10)

    def forward(self, x):
        x = self.conv1(x)
        x = self.conv2(x)
        x = x.view(x.size(0), -1)
        output = self.linear(x)
        return output

net = Net()
print(net)

運行結果：

Net(
  (conv1): Sequential(
    (conv1): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (ReLU1): ReLU()
    (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (conv2): Sequential(
    (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (ReLU2): ReLU()
    (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (linear): Linear(in_features=1568, out_features=10, bias=True)
)

torch.nn.Modules.children()

children()這個屬性返回下一級模塊的迭代器(可與下一小節中的modules()模塊對照，更容易理解)。

i = 1
for layer in net.children():
    print('{}th layer:'.format(i))
    print(layer)
    i += 1

運行結果：

1th layer:
Sequential(
  (conv1): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (ReLU1): ReLU()
  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
2th layer:
Sequential(
  (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (ReLU2): ReLU()
  (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
3th layer:
Linear(in_features=1568, out_features=10, bias=True)

torch.nn.Modules.modules()

 modules()屬性返回神經網絡每一級的內容。

for layer in net.modules():
    print(layer)
    print()

運行結果：

Net(
  (conv1): Sequential(
    (conv1): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (ReLU1): ReLU()
    (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (conv2): Sequential(
    (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (ReLU2): ReLU()
    (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (linear): Linear(in_features=1568, out_features=10, bias=True)
)

Sequential(
  (conv1): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (ReLU1): ReLU()
  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)

Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))

ReLU()

MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)

Sequential(
  (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (ReLU2): ReLU()
  (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)

Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))

ReLU()

MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)

Linear(in_features=1568, out_features=10, bias=True)

從上面的運行結果可以看出來，modules()模塊以迭代器的形式返回神經網絡逐級的內容。

torch.nn.Modules.named_children()和torch.nn.Modules.named_modules()

named_children()和named_modules()這兩個模塊不僅返回模塊的迭代器，而且還會返回網絡層的名字。

for layer in net.named_modules():
    print(layer)

運行結果：

('', Net(
  (conv1): Sequential(
    (conv1): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (ReLU1): ReLU()
    (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (conv2): Sequential(
    (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
    (ReLU2): ReLU()
    (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (linear): Linear(in_features=1568, out_features=10, bias=True)
))
('conv1', Sequential(
  (conv1): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (ReLU1): ReLU()
  (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
))
('conv1.conv1', Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2)))
('conv1.ReLU1', ReLU())
('conv1.pool1', MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False))
('conv2', Sequential(
  (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (ReLU2): ReLU()
  (pool2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
))
('conv2.conv2', Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2)))
('conv2.ReLU2', ReLU())
('conv2.pool2', MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False))
('linear', Linear(in_features=1568, out_features=10, bias=True))

只提取模型中的卷積層

如果希望只提取卷積層，可以這樣操作：

conv_net = nn.Sequential()
for layer in net.named_modules():
    if isinstance(layer[1], nn.Conv2d):
        conv_net.add_module(layer[0][:5], layer[1])    # layer[1] is the net layer, layer[0] is its name.
print(conv_net)

運行結果：

Sequential(
  (conv1): Conv2d(1, 16, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
  (conv2): Conv2d(16, 32, kernel_size=(5, 5), stride=(1, 1), padding=(2, 2))
)

參數初始化

參數初始化需用到torch.nn.init，下面幾個函數比較常用，截圖摘自PyTorch中文文檔。

在神經網絡中，模型參數就是weights和bias，weights和bias在PyTorch中是Variable，所以我們只需取出weights的data屬性，並填充數值就可以了。比如對卷積層進行參數初始化：

for m in net.modules():
    if isinstance(m, nn.Conv2d):
        nn.init.normal(m.weight.data)
        # nn.init.xavier_normal(m.weight.data)
        # nn.init.kaiming_normal(m.weight.data)
    elif isinstance(m, nn.Linear):
        m.weight.data.normal_()
    m.bias.data.fill_(0)

 

參考文獻：

深度學習之PyTorch第四章，廖星宇.