PyTorch中nn.Module类简介
? ? ? torch.nn.Module類是所有神經(jīng)網(wǎng)絡(luò)模塊(modules)的基類，它的實現(xiàn)在torch/nn/modules/module.py中。你的模型也應(yīng)該繼承這個類，主要重載__init__、forward和extra_repr函數(shù)。Modules還可以包含其它Modules，從而可以將它們嵌套在樹結(jié)構(gòu)中。
? ? ? 只要在自己的類中定義了forward函數(shù)，backward函數(shù)就會利用Autograd被自動實現(xiàn)。只要實例化一個對象并傳入對應(yīng)的參數(shù)就可以自動調(diào)用forward函數(shù)。因為此時會調(diào)用對象的__call__方法，而nn.Module類中的__call__方法會調(diào)用forward函數(shù)。
? ? ? nn.Module類中函數(shù)介紹：
? ? ? __init__：初始化內(nèi)部module狀態(tài)。
? ? ? register_buffer：向module添加buffer，不作為模型參數(shù)，可作為module狀態(tài)的一部分。默認(rèn)情況下，buffer是持久(persistent)的，將與參數(shù)一起保存。buffer是否persistent的區(qū)別在于這個buffer是否被放入self.state_dict()中被保存下來。
? ? ? register_parameter：向module添加參數(shù)。
? ? ? add_module：添加一個submodule(children)到當(dāng)前module中。
? ? ? apply：將fn遞歸應(yīng)用于每個submodule(children)，典型用途為初始化模型參數(shù)。
? ? ? cuda：將所有模型參數(shù)和buffers轉(zhuǎn)移到GPU上。
? ? ? xpu：將所有模型參數(shù)和buffers轉(zhuǎn)移到XPU上。
? ? ? cpu：將所有模型參數(shù)和buffers轉(zhuǎn)移到CPU上。
? ? ? type：將所有參數(shù)和buffers轉(zhuǎn)換為所需的類型。
? ? ? float：將所有浮點參數(shù)和buffers轉(zhuǎn)換為float32數(shù)據(jù)類型。
? ? ? double：將所有浮點參數(shù)和buffers轉(zhuǎn)換為double數(shù)據(jù)類型。
? ? ? half：將所有浮點參數(shù)和buffers轉(zhuǎn)換為float16數(shù)據(jù)類型。
? ? ? bfloat16：將所有浮點參數(shù)和buffers轉(zhuǎn)換為bfloat16數(shù)據(jù)類型。
? ? ? to：將參數(shù)和buffers轉(zhuǎn)換為指定的數(shù)據(jù)類型或轉(zhuǎn)換到指定的設(shè)備上。
? ? ? register_backward_hook：在module中注冊一個反向鉤子。不推薦使用。
? ? ? register_full_backward_hook：在module中注冊一個反向鉤子。每次計算梯度時都會調(diào)用此鉤子。使用此鉤子時不允許就地(in place)修改輸入或輸出，否則會觸發(fā)error。
? ? ? register_forward_pre_hook：在module中注冊前向pre-hook。每次調(diào)用forward之前都會調(diào)用此鉤子。
? ? ? register_forward_hook：在module中注冊一個前向鉤子。每次forward計算輸出后都會調(diào)用此鉤子。
? ? ? state_dict：返回包含了module的整個狀態(tài)的字典。其中keys是對應(yīng)的參數(shù)和buffer名稱。
? ? ? load_state_dict：將參數(shù)和buffers從state_dict復(fù)制到module及其后代(descendants)中。
? ? ? parameters：返回module的參數(shù)的迭代器。
? ? ? named_parameters：返回module的參數(shù)的迭代器，產(chǎn)生(yield)參數(shù)的名稱以及參數(shù)本身。不會返回重復(fù)的parameter。
? ? ? buffers：返回module的buffers的迭代器。
? ? ? named_buffers：返回module的buffers的迭代器，產(chǎn)生(yield)buffer的名稱以及buffer本身。不會返回重復(fù)的buffer。
? ? ? children：返回直接子module的迭代器。
? ? ? named_children：返回直接子module的迭代器，產(chǎn)生(yield)子module的名稱以及子module本身。不會返回重復(fù)的children。
? ? ? modules：返回網(wǎng)絡(luò)中所有modules的迭代器。
? ? ? named_modules：返回網(wǎng)絡(luò)中所有modules的迭代器，產(chǎn)生(yield)module的名稱以及module本身。不會返回重復(fù)的module。
? ? ? train：將module設(shè)置為訓(xùn)練模式。這僅對某些module起作用。module.py實現(xiàn)中會修改self.training并通過self.children()來調(diào)整所有submodule的狀態(tài)。
? ? ? eval：將module設(shè)置為評估模式。這僅對某些module起作用。module.py實現(xiàn)中直接調(diào)用train(False)。
? ? ? requires_grad_：更改autograd是否應(yīng)記錄對此module中參數(shù)的操作。此方法就地(in place)設(shè)置參數(shù)的requires_grad屬性。
? ? ? zero_grad：將所有模型參數(shù)的梯度設(shè)置為零。
? ? ? share_memory：
? ? ? extra_repr：設(shè)置module的額外表示。你應(yīng)該在自己的modules中重新實現(xiàn)此方法。
? ? ?測試代碼如下：
import torch
import torch.nn as nn
import torch.nn.functional as F # nn.functional.py中存放激活函數(shù)等的實現(xiàn)@torch.no_grad()
def init_weights(m):print("xxxx:", m)if type(m) == nn.Linear:m.weight.fill_(1.0)print("yyyy:", m.weight)class Model(nn.Module):def __init__(self):# 在實現(xiàn)自己的__init__函數(shù)時,為了正確初始化自定義的神經(jīng)網(wǎng)絡(luò)模塊,一定要先調(diào)用super().__init__super(Model, self).__init__()self.conv1 = nn.Conv2d(1, 20, 5) # submodule(child module)self.conv2 = nn.Conv2d(20, 20, 5)self.add_module("conv3", nn.Conv2d(10, 40, 5)) # 添加一個submodule到當(dāng)前module,等價于self.conv3 = nn.Conv2d(10, 40, 5)self.register_buffer("buffer", torch.randn([2,3])) # 給module添加一個presistent(持久的) bufferself.param1 = nn.Parameter(torch.rand([1])) # module參數(shù)的tensorself.register_parameter("param2", nn.Parameter(torch.rand([1]))) # 向module添加參數(shù)# nn.Sequential: 順序容器,module將按照它們在構(gòu)造函數(shù)中傳遞的順序添加,它允許將整個容器視為單個moduleself.feature = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))self.feature.apply(init_weights) # 將fn遞歸應(yīng)用于每個submodule,典型用途為初始化模型參數(shù)self.feature.to(torch.double) # 將參數(shù)數(shù)據(jù)類型轉(zhuǎn)換為doublecpu = torch.device("cpu")self.feature.to(cpu) # 將參數(shù)數(shù)據(jù)轉(zhuǎn)換到cpu設(shè)備上def forward(self, x):x = F.relu(self.conv1(x))return F.relu(self.conv2(x))model = Model()
print("## Model:", model)model.cpu() # 將所有模型參數(shù)和buffers移動到CPU上
model.float() # 將所有浮點參數(shù)和buffers轉(zhuǎn)換為float數(shù)據(jù)類型
model.zero_grad() # 將所有模型參數(shù)的梯度設(shè)置為零# state_dict:返回一個字典,保存著module的所有狀態(tài),參數(shù)和persistent buffers都會包含在字典中,字典的key就是參數(shù)和buffer的names
print("## state_dict:", model.state_dict().keys())for name, parameters in model.named_parameters(): # 返回module的參數(shù)(weight and bias)的迭代器,產(chǎn)生(yield)參數(shù)的名稱以及參數(shù)本身print(f"## named_parameters: name: {name}; parameters size: {parameters.size()}")for name, buffers in model.named_buffers(): # 返回module的buffers的迭代器,產(chǎn)生(yield)buffer的名稱以及buffer本身print(f"## named_buffers: name: {name}; buffers size: {buffers.size()}")# 注:children和modules中重復(fù)的module只被返回一次
for children in model.children(): # 返回當(dāng)前module的child module(submodule)的迭代器print("## children:", children)for name, children in model.named_children(): # 返回直接submodule的迭代器,產(chǎn)生(yield) submodule的名稱以及submodule本身print(f"## named_children: name: {name}; children: {children}")for modules in model.modules(): # 返回當(dāng)前模型所有module的迭代器,注意與children的區(qū)別print("## modules:", modules)for name, modules in model.named_modules(): # 返回網(wǎng)絡(luò)中所有modules的迭代器,產(chǎn)生(yield)module的名稱以及module本身,注意與named_children的區(qū)別print(f"## named_modules: name: {name}; module: {modules}")model.train() # 將module設(shè)置為訓(xùn)練模式
model.eval() # 將module設(shè)置為評估模式print("test finish")? ? ?GitHub：https://github.com/fengbingchun/PyTorch_Test
總結(jié)
以上是生活随笔為你收集整理的PyTorch中nn.Module类简介的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。
- 上一篇： Python3中装饰器@typing.o
- 下一篇： Python3中参数*args和**kw