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第四章，廖星宇.