关于PyTorch中nn.Module类的简介
作者:fengbingchun
这篇文章主要介绍了关于PyTorch中nn.Module类的简介,具有很好的参考价值,希望对大家有所帮助。如有错误或未考虑完全的地方,望不吝赐教
PyTorch nn.Module类的简介
torch.nn.Module类是所有神经网络模块(modules)的基类,它的实现在torch/nn/modules/module.py中。你的模型也应该继承这个类,主要重载__init__、forward和extra_repr函数。Modules还可以包含其它Modules,从而可以将它们嵌套在树结构中。
只要在自己的类中定义了forward函数,backward函数就会利用Autograd被自动实现。只要实例化一个对象并传入对应的参数就可以自动调用forward函数。因为此时会调用对象的__call__方法,而nn.Module类中的__call__方法会调用forward函数。
nn.Module类中函数介绍:
__init__
:初始化内部module状态。register_buffer
:向module添加buffer,不作为模型参数,可作为module状态的一部分。默认情况下,buffer是持久(persistent)的,将与参数一起保存。buffer是否persistent的区别在于这个buffer是否被放入self.state_dict()中被保存下来。register_parameter
:向module添加参数。add_module
:添加一个submodule(children)到当前module中。apply
:将fn递归应用于每个submodule(children),典型用途为初始化模型参数。cuda
:将所有模型参数和buffers转移到GPU上。xpu
:将所有模型参数和buffers转移到XPU上。cpu
:将所有模型参数和buffers转移到CPU上。type
:将所有参数和buffers转换为所需的类型。float
:将所有浮点参数和buffers转换为float32数据类型。double
:将所有浮点参数和buffers转换为double数据类型。half
:将所有浮点参数和buffers转换为float16数据类型。bfloat16
:将所有浮点参数和buffers转换为bfloat16数据类型。to
:将参数和buffers转换为指定的数据类型或转换到指定的设备上。register_backward_hook
:在module中注册一个反向钩子。不推荐使用。register_full_backward_hook
:在module中注册一个反向钩子。每次计算梯度时都会调用此钩子。使用此钩子时不允许就地(in place)修改输入或输出,否则会触发error。register_forward_pre_hook
:在module中注册前向pre-hook。每次调用forward之前都会调用此钩子。register_forward_hook
:在module中注册一个前向钩子。每次forward计算输出后都会调用此钩子。state_dict
:返回包含了module的整个状态的字典。其中keys是对应的参数和buffer名称。load_state_dict
:将参数和buffers从state_dict复制到module及其后代(descendants)中。parameters
:返回module的参数的迭代器。named_parameters
:返回module的参数的迭代器,产生(yield)参数的名称以及参数本身。不会返回重复的parameter。buffers
:返回module的buffers的迭代器。named_buffers
:返回module的buffers的迭代器,产生(yield)buffer的名称以及buffer本身。不会返回重复的buffer。children
:返回直接子module的迭代器。named_children
:返回直接子module的迭代器,产生(yield)子module的名称以及子module本身。不会返回重复的children。modules
:返回网络中所有modules的迭代器。named_modules
:返回网络中所有modules的迭代器,产生(yield)module的名称以及module本身。不会返回重复的module。train
:将module设置为训练模式。这仅对某些module起作用。module.py实现中会修改self.training并通过self.children()来调整所有submodule的状态。eval
:将module设置为评估模式。这仅对某些module起作用。module.py实现中直接调用train(False)。requires_grad
_:更改autograd是否应记录对此module中参数的操作。此方法就地(in place)设置参数的requires_grad属性。zero_grad
:将所有模型参数的梯度设置为零。extra_repr
:设置module的额外表示。你应该在自己的modules中重新实现此方法。
测试代码如下:
import torch import torch.nn as nn import torch.nn.functional as F # nn.functional.py中存放激活函数等的实现 @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): # 在实现自己的__init__函数时,为了正确初始化自定义的神经网络模块,一定要先调用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到当前module,等价于self.conv3 = nn.Conv2d(10, 40, 5) self.register_buffer("buffer", torch.randn([2,3])) # 给module添加一个presistent(持久的) buffer self.param1 = nn.Parameter(torch.rand([1])) # module参数的tensor self.register_parameter("param2", nn.Parameter(torch.rand([1]))) # 向module添加参数 # nn.Sequential: 顺序容器,module将按照它们在构造函数中传递的顺序添加,它允许将整个容器视为单个module self.feature = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2)) self.feature.apply(init_weights) # 将fn递归应用于每个submodule,典型用途为初始化模型参数 self.feature.to(torch.double) # 将参数数据类型转换为double cpu = torch.device("cpu") self.feature.to(cpu) # 将参数数据转换到cpu设备上 def forward(self, x): x = F.relu(self.conv1(x)) return F.relu(self.conv2(x)) model = Model() print("## Model:", model) model.cpu() # 将所有模型参数和buffers移动到CPU上 model.float() # 将所有浮点参数和buffers转换为float数据类型 model.zero_grad() # 将所有模型参数的梯度设置为零 # state_dict:返回一个字典,保存着module的所有状态,参数和persistent buffers都会包含在字典中,字典的key就是参数和buffer的names print("## state_dict:", model.state_dict().keys()) for name, parameters in model.named_parameters(): # 返回module的参数(weight and bias)的迭代器,产生(yield)参数的名称以及参数本身 print(f"## named_parameters: name: {name}; parameters size: {parameters.size()}") for name, buffers in model.named_buffers(): # 返回module的buffers的迭代器,产生(yield)buffer的名称以及buffer本身 print(f"## named_buffers: name: {name}; buffers size: {buffers.size()}") # 注:children和modules中重复的module只被返回一次 for children in model.children(): # 返回当前module的child module(submodule)的迭代器 print("## children:", children) for name, children in model.named_children(): # 返回直接submodule的迭代器,产生(yield) submodule的名称以及submodule本身 print(f"## named_children: name: {name}; children: {children}") for modules in model.modules(): # 返回当前模型所有module的迭代器,注意与children的区别 print("## modules:", modules) for name, modules in model.named_modules(): # 返回网络中所有modules的迭代器,产生(yield)module的名称以及module本身,注意与named_children的区别 print(f"## named_modules: name: {name}; module: {modules}") model.train() # 将module设置为训练模式 model.eval() # 将module设置为评估模式 print("test finish")
GitHub:https://github.com/fengbingchun/PyTorch_Test
PyTorch中nn.Module理解
nn.Module是Pytorch封装的一个类,是搭建神经网络时需要继承的父类:
import torch import torch.nn as nn # 括号中加入nn.Module(父类)。Test2变成子类,继承父类(nn.Module)的所有特性。 class Test2(nn.Module): def __init__(self): # Test2类定义初始化方法 super(Test2, self).__init__() # 父类初始化 self.M = nn.Parameter(torch.ones(10)) def weightInit(self): print('Testing') def forward(self, n): # print(2 * n) print(self.M * n) self.weightInit() # 调用方法 network = Test2() network(2) # 2赋值给forward(self, n)中的n。 ……省略一部分代码…… # 因为Test2是nn.Module的子类,所以也可以执行父类中的方法。如: model_dict = network.state_dict() # 调用父类中的方法state_dict(),将Test2中训练参数赋值model_dict。 for k, v in model_dict.items(): # 查看自己网络参数各层名称、数值 print(k) # 输出网络参数名字 # print(v) # 输出网络参数数值
继承nn.Module的子类程序是从forward()方法开始执行的,如果要想执行其他方法,必须把它放在forward()方法中。这一点与python中继承有稍许的不同。
总结
以上为个人经验,希望能给大家一个参考,也希望大家多多支持脚本之家。