之前有一篇文章介绍了使用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
----------------------------------------------------------------