pytorch中使用add_module添加网络子模块

之前有一篇文章介绍了使用Sequential、ModuleList和ModuleDict添加网络，除此之外，我们还可以通过add_module()添加每一层，并且为每一层增加了一个单独的名字。add_module()可以快速地替换特定结构可以不用修改过多的代码。

add_module的功能为Module添加一个子module，对应名字为name。使用方式如下：

add_module(name, module)

其中name为子模块的名字，使用这个名字可以访问特定的子module，module为我们自定义的子module。

一般情况下子module都是在A.init(self)中定义的，比如A中一个卷积子模块self.conv1 = torch.nn.Conv2d(…)。此时，这个卷积模块在A的名字其实是’conv1’。

对比之下，add_module()函数就可以在A.init(self)以外定义A的子模块。如定义同样的卷积子模块，可以通过A.add_module(‘conv1’, torch.nn.Conv2d(…))。

直接看代码会更直观，举例如下：

• 普通add_module举例

from torch import nn
from torchsummary import summary

class Net1(nn.Module):
    def __init__(self):
        super(Net1,self).__init__()
        self.conv_1 = nn.Conv2d(3,6,3)
        self.add_module('conv_2', nn.Conv2d(6,12,3))
        self.conv_3 = nn.Conv2d(12,24,3)
        
    def forward(self,x):
          x = self.conv_1(x)
          x = self.conv_2(x) # 使用name访问
          x = self.conv_3(x)
          return x
    
model = Net1()
print(model)
model.to('cuda')
summary(model,(3,128,128)) 

输出：

Net1(
  (conv_dd): Conv2d(3, 6, kernel_size=(3, 3), stride=(1, 1))
  (conv_2): Conv2d(6, 12, kernel_size=(3, 3), stride=(1, 1))
  (conv_3): Conv2d(12, 24, kernel_size=(3, 3), stride=(1, 1))
)
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1          [-1, 6, 126, 126]             168
            Conv2d-2         [-1, 12, 124, 124]             660
            Conv2d-3         [-1, 24, 122, 122]           2,616
================================================================
Total params: 3,444
Trainable params: 3,444
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.19
Forward/backward pass size (MB): 4.86
Params size (MB): 0.01
Estimated Total Size (MB): 5.06
----------------------------------------------------------------

• 使用add_module添加至Sequential模块举例

class Net3(torch.nn.Module):
  def __init__(self):
    super(Net3, self).__init__()
    self.conv=torch.nn.Sequential()
    self.conv.add_module("conv1",torch.nn.Conv2d(3, 32, 3, 1, 1)) # 没有下采样，kernel = 3, stride = 1, padding = 1
    self.conv.add_module("relu1",torch.nn.ReLU())
    self.conv.add_module("pool1",torch.nn.MaxPool2d(2)) # 下采样，fearture size减半
    self.dense = torch.nn.Sequential()
    self.dense.add_module("dense1",torch.nn.Linear(32 * 64 * 64, 128)) # 输入cxhxw, 需要对应上面的输出
    self.dense.add_module("relu2",torch.nn.ReLU())
    self.dense.add_module("dense2",torch.nn.Linear(128, 10))
 
  def forward(self, x):
        conv_out = self.conv(x)
        res = conv_out.view(conv_out.size(0), -1) # 重新排列，进行线性运算
        out = self.dense(res)
        return out
 
model3 = Net3()
print(model3)

print(model3)
model3.to('cuda')
summary(model3,(3,128,128))

输出如下：

Net3(
  (conv): Sequential(
    (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (relu1): ReLU()
    (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (dense): Sequential(
    (dense1): Linear(in_features=131072, out_features=128, bias=True)
    (relu2): ReLU()
    (dense2): Linear(in_features=128, out_features=10, bias=True)
  )
)
Net3(
  (conv): Sequential(
    (conv1): Conv2d(3, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (relu1): ReLU()
    (pool1): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (dense): Sequential(
    (dense1): Linear(in_features=131072, out_features=128, bias=True)
    (relu2): ReLU()
    (dense2): Linear(in_features=128, out_features=10, bias=True)
  )
)
----------------------------------------------------------------
        Layer (type)               Output Shape         Param #
================================================================
            Conv2d-1         [-1, 32, 128, 128]             896
              ReLU-2         [-1, 32, 128, 128]               0
         MaxPool2d-3           [-1, 32, 64, 64]               0
            Linear-4                  [-1, 128]      16,777,344
              ReLU-5                  [-1, 128]               0
            Linear-6                   [-1, 10]           1,290
================================================================
Total params: 16,779,530
Trainable params: 16,779,530
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.19
Forward/backward pass size (MB): 9.00
Params size (MB): 64.01
Estimated Total Size (MB): 73.20
----------------------------------------------------------------