ModuleList
是特殊的list,其包含的模块会被自动注册,对所有的Module
方法都可见。先给结论:如果要用列表组织模型模块,那么强烈建议使用nn.ModuleList。这有什么好处呢?看下面的例子。
import torch.nn as nn
from torchsummary import summary
class MyNet(nn.Module):
def __init__(self):
super().__init__()
self.net = [
nn.Conv2d(3, 64, 3, padding='same'),
nn.Conv2d(64, 64, 3, padding='same'),
nn.Conv2d(64, 64, 3, padding='same')
]
self.conv = nn.Conv2d(64, 64, 3, padding='same')
def forward(self, x):
for net_ in self.net:
x = net_(x)
out = self.conv(x)
return x
model = MyNet()
print(summary(model, (3, 224, 224)))
MyNet
的部分模块包含在list
中,使用torchsummary
模块中的summary
方法查看模型结构,得到输出如下:
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 224, 224] 36,928
================================================================
Total params: 36,928
Trainable params: 36,928
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 24.50
Params size (MB): 0.14
Estimated Total Size (MB): 25.22
----------------------------------------------------------------
可以看到,包含在list中的模块无法被展示出来,而将list替换为ModuleList则可以解决这个问题。
就这?就这??这有啥用,我不用torchsummary
这个模块不就行了[doge]!
当然不是。torchsummary.summary()
方法失效只是问题现象,根本原因在于list中的模块未能和模型很紧密的绑定在一起,模型压根就不知道这里面模块的存在!当使用model.to(device)
时会带来很严重的错误。看下面这段代码,由于网络不知道list中模块的存在,所以当移动模型位置(例如从cpu到gpu时),list中的模块会被略掉,这样当前向传播计算时会因为Tensor不在同一个设备上而出错。
print(model.to('cpu'))
"""output
MyNet(
(conv): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=same)
)
"""