快捷導(dǎo)航

Pytorch參數(shù)注冊(cè)和nn.ModuleList nn.ModuleDict的問(wèn)題

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

這篇文章主要介紹了Pytorch參數(shù)注冊(cè)和nn.ModuleList nn.ModuleDict的問(wèn)題，具有很好的參考價(jià)值，希望對(duì)大家有所幫助。如有錯(cuò)誤或未考慮完全的地方，望不吝賜教

參數(shù)注冊(cè)

嘗試自己寫GoogLeNet時(shí)碰到的問(wèn)題，放在字典中的參數(shù)無(wú)法自動(dòng)注冊(cè)，所謂的注冊(cè)，就是當(dāng)參數(shù)注冊(cè)到這個(gè)網(wǎng)絡(luò)上時(shí)，它會(huì)隨著你在外部調(diào)用net.cuda()后自動(dòng)遷移到GPU上，而沒(méi)有注冊(cè)的參數(shù)則不會(huì)隨著網(wǎng)絡(luò)遷到GPU上，這就可能導(dǎo)致輸入在GPU上而參數(shù)不在GPU上，從而出現(xiàn)錯(cuò)誤，為了說(shuō)明這個(gè)現(xiàn)象。

舉一個(gè)有點(diǎn)鐵憨憨的例子：

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)) # 這里其實(shí)可以直接用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))

上面的代碼可以成功運(yùn)行，因?yàn)樗械臄?shù)值都是放在CPU上的，但是，一旦我們要把模型移到GPU上時(shí)

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

運(yùn)行后就會(huì)出現(xiàn)

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

這就是因?yàn)閟elf.weight沒(méi)有隨著模型一起移到GPU上的原因，此時(shí)我們查看模型的參數(shù)，會(huì)發(fā)現(xiàn)并沒(méi)有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)

上面的代碼沒(méi)有輸出，因?yàn)閚et根本沒(méi)有參數(shù)

那么為了讓net有參數(shù)，我們需要手動(dòng)地將self.weight注冊(cè)到網(wǎng)絡(luò)上

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))) # 被注冊(cè)的參數(shù)必須是nn.Parameter類型
?? ??? ?self.register_parameter('weight',self.weight) # 手動(dòng)注冊(cè)參數(shù)
?? ??? ?
?? ?
?? ?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))

此時(shí)網(wǎng)絡(luò)的參數(shù)就有了輸出，同時(shí)會(huì)隨著一起遷到GPU上，輸出就類似這樣

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

不過(guò)后來(lái)我實(shí)驗(yàn)了以下，好像只寫nn.Parameter不寫register也可以被默認(rèn)注冊(cè)

nn.ModuleList和nn.ModuleDict

有時(shí)候我們?yōu)榱藞D省事，可能會(huì)這樣寫網(wǎng)絡(luò)

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

同樣，輸出網(wǎng)絡(luò)的參數(shù)啥也沒(méi)有，這意味著當(dāng)調(diào)用net.cuda時(shí)，self.linears里面的參數(shù)不會(huì)一起走到GPU上去

此時(shí)我們可以在__init__方法中手動(dòng)對(duì)self.parameters()迭代然后把每個(gè)參數(shù)注冊(cè)，但更好的方法是，pytorch已經(jīng)為我們提供了nn.ModuleList，用來(lái)代替python內(nèi)置的list，放在nn.ModuleList中的參數(shù)將會(huì)自動(dòng)被正確注冊(cè)

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

此時(shí)就有輸出了

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

nn.ModuleDict也是類似，當(dāng)我們需要把參數(shù)放在一個(gè)字典里的時(shí)候，能夠用的上，這里直接給一個(gè)官方的例子看一看就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中的參數(shù)不會(huì)自動(dòng)注冊(cè)，但如果只是暫時(shí)放在list里，隨后又調(diào)用了nn.Sequential把整個(gè)list整合起來(lái)，參數(shù)仍然是會(huì)自動(dòng)注冊(cè)的

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