利用add_module ,添加conv_2.
add_module用法:add_module(name,module)
直接看代码
from torch import nn
from torchsummary import summary
class Net_test(nn.Module):
def __init__(self):
super(Net_test,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)
x = self.conv_3(x)
return x
model = Net_test()
print(model)
model.to('cuda')
summary(model,(3,128,128))
结果
Net_test(
(conv_1): 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
----------------------------------------------------------------
可以看到 conv_2 已经成功的加入到网络中