`nn.ModuleList`和普通list的區別


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)
)
"""


免責聲明!

本站轉載的文章為個人學習借鑒使用,本站對版權不負任何法律責任。如果侵犯了您的隱私權益,請聯系本站郵箱yoyou2525@163.com刪除。



 
粵ICP備18138465號   © 2018-2025 CODEPRJ.COM