Pytorch參數注冊和nn.ModuleList nn.ModuleDict的問題

更新時間：2023年01月03日 09:22:39 作者：luputo

這篇文章主要介紹了Pytorch參數注冊和nn.ModuleList nn.ModuleDict的問題，具有很好的參考價值，希望對大家有所幫助。如有錯誤或未考慮完全的地方，望不吝賜教

參數注冊

嘗試自己寫GoogLeNet時碰到的問題，放在字典中的參數無法自動注冊，所謂的注冊，就是當參數注冊到這個網絡上時，它會隨著你在外部調用net.cuda()后自動遷移到GPU上，而沒有注冊的參數則不會隨著網絡遷到GPU上，這就可能導致輸入在GPU上而參數不在GPU上，從而出現錯誤，為了說明這個現象。

舉一個有點鐵憨憨的例子：

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
?? ?def __init__(self):
?? ??? ?super(Net,self).__init__()
?? ??? ?self.weight = torch.rand((3,4)) # 這里其實可以直接用nn.Linear，但為了舉例這里先憨憨一下
?? ?
?? ?def forward(self,x):
?? ??? ?return F.linear(x,self.weight)

if __name__ == "__main__":
?? ?batch_size = 10
?? ?dummy = torch.rand((batch_size,4))
?? ?net = Net()
?? ?print(net(dummy))

上面的代碼可以成功運行，因為所有的數值都是放在CPU上的，但是，一旦我們要把模型移到GPU上時

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
?? ?def __init__(self):
?? ??? ?super(Net,self).__init__()
?? ??? ?self.weight = torch.rand((3,4))
?? ?
?? ?def forward(self,x):
?? ??? ?return F.linear(x,self.weight)

if __name__ == "__main__":
?? ?batch_size = 10
?? ?dummy = torch.rand((batch_size,4)).cuda()
?? ?net = Net().cuda()
?? ?print(net(dummy))

運行后就會出現

...
RuntimeError: Expected object of backend CUDA but got backend CPU for argument #2 'mat2'

這就是因為self.weight沒有隨著模型一起移到GPU上的原因，此時我們查看模型的參數，會發(fā)現并沒有self.weight

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
?? ?def __init__(self):
?? ??? ?super(Net,self).__init__()
?? ??? ?self.weight = torch.rand((3,4))
?? ?
?? ?def forward(self,x):
?? ??? ?return F.linear(x,self.weight)

if __name__ == "__main__":
?? ?net = Net()
?? ?for parameter in net.parameters():
?? ??? ?print(parameter)

上面的代碼沒有輸出，因為net根本沒有參數

那么為了讓net有參數，我們需要手動地將self.weight注冊到網絡上

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
?? ?def __init__(self):
?? ??? ?super(Net,self).__init__()
?? ??? ?self.weight = nn.Parameter(torch.rand((3,4))) # 被注冊的參數必須是nn.Parameter類型
?? ??? ?self.register_parameter('weight',self.weight) # 手動注冊參數
?? ??? ?
?? ?
?? ?def forward(self,x):
?? ??? ?return F.linear(x,self.weight)

if __name__ == "__main__":
?? ?net = Net()
?? ?for parameter in net.parameters():
?? ??? ?print(parameter)

?? ?batch_size = 10
?? ?net = net.cuda()
?? ?dummy = torch.rand((batch_size,4)).cuda()
?? ?print(net(dummy))

此時網絡的參數就有了輸出，同時會隨著一起遷到GPU上，輸出就類似這樣

Parameter containing:
tensor([...])
tensor([...])

不過后來我實驗了以下，好像只寫nn.Parameter不寫register也可以被默認注冊

nn.ModuleList和nn.ModuleDict

有時候我們?yōu)榱藞D省事，可能會這樣寫網絡

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
?? ?def __init__(self):
?? ??? ?super(Net,self).__init__()
?? ??? ?self.linears = [nn.Linear(4,4),nn.Linear(4,4),nn.Linear(4,2)]
?? ?
?? ?def forward(self,x):
?? ??? ?for linear in self.linears:
?? ??? ??? ?x = linear(x)
?? ??? ??? ?x = F.relu(x)
?? ??? ?return x

if __name__ == '__main__':
?? ?net = Net()
?? ?for parameter in net.parameters():
?? ??? ?print(parameter)??

同樣，輸出網絡的參數啥也沒有，這意味著當調用net.cuda時，self.linears里面的參數不會一起走到GPU上去

此時我們可以在__init__方法中手動對self.parameters()迭代然后把每個參數注冊，但更好的方法是，pytorch已經為我們提供了nn.ModuleList，用來代替python內置的list，放在nn.ModuleList中的參數將會自動被正確注冊

import torch
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
?? ?def __init__(self):
?? ??? ?super(Net,self).__init__()
?? ??? ?self.linears = nn.ModuleList([nn.Linear(4,4),nn.Linear(4,4),nn.Linear(4,2)])
?? ?
?? ?def forward(self,x):
?? ??? ?for linear in self.linears:
?? ??? ??? ?x = linear(x)
?? ??? ??? ?x = F.relu(x)
?? ??? ?return x

if __name__ == '__main__':
?? ?net = Net()
?? ?for parameter in net.parameters():
?? ??? ?print(parameter)?? ??? ?

此時就有輸出了

Parameter containing:
tensor(...)
Parameter containing:
tensor(...)
...

nn.ModuleDict也是類似，當我們需要把參數放在一個字典里的時候，能夠用的上，這里直接給一個官方的例子看一看就OK

class MyModule(nn.Module):
? ? def __init__(self):
? ? ? ? super(MyModule, self).__init__()
? ? ? ? self.choices = nn.ModuleDict({
? ? ? ? ? ? ? ? 'conv': nn.Conv2d(10, 10, 3),
? ? ? ? ? ? ? ? 'pool': nn.MaxPool2d(3)
? ? ? ? })
? ? ? ? self.activations = nn.ModuleDict([
? ? ? ? ? ? ? ? ['lrelu', nn.LeakyReLU()],
? ? ? ? ? ? ? ? ['prelu', nn.PReLU()]
? ? ? ? ])

? ? def forward(self, x, choice, act):
? ? ? ? x = self.choices[choice](x)
? ? ? ? x = self.activations[act](x)
? ? ? ? return x

需要注意的是，雖然直接放在python list中的參數不會自動注冊，但如果只是暫時放在list里，隨后又調用了nn.Sequential把整個list整合起來，參數仍然是會自動注冊的

另外一點要注意的是ModuleList和ModuleDict里面只能放Module的子類，也就是nn.Conv,nn.Linear這樣的，但不能放nn.Parameter，如果要放nn.Parameter，用nn.ParameterList即可，用法和nn.ModuleList一樣