当前位置：首页 > 编程日记 > 正文

PyTorch中nn.Module类简介

编程日记 2024-12-13 09:40:00

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：将所有模型参数的梯度设置为零。

share_memory：

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(持久的) bufferself.param1 = nn.Parameter(torch.rand([1])) # module参数的tensorself.register_parameter("param2", nn.Parameter(torch.rand([1]))) # 向module添加参数# nn.Sequential: 顺序容器,module将按照它们在构造函数中传递的顺序添加,它允许将整个容器视为单个moduleself.feature = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))self.feature.apply(init_weights) # 将fn递归应用于每个submodule,典型用途为初始化模型参数self.feature.to(torch.double) # 将参数数据类型转换为doublecpu = 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

https://www.dkcj.cn/info/31628.html

PyTorch中nn.Module类简介

相关文章：

什么是三层交换机、网关、DNS、子网掩码、MAC地址

20行代码发一篇NeurIPS：梯度共享已经不安全了

关于页面打印window.print()的样式问题

Python3中参数*args和**kwargs介绍

4大主流CPU处理器技术架构，不知道就out了！

grunt-connect-proxy解决开发时跨域问题

混合推荐系统就是多个推荐系统“大杂烩”吗？

Python3中collections.OrderedDict介绍

汗！雅虎中国个人空间

关于v$process与v$session中process的理解

Python3中lambda表达式介绍

6大理由，告诉你为什么这个大会你不能错过！ | 文末有福利

魔与道的反复较量反垃圾邮件技术

在Centos 7下编译openwrt+njit-client

Python3中内置函数callable介绍

户外广告新创意

百度重新定义「智能屏」，瞄准10后

jQuery最简单的表单提交方式

Python3中typing模块介绍

显示所有文件和文件夹无论如何无法被设置

工作5年后才明白的道理：不起眼的技能中，藏着你的未来

Transform-style和Perspective属性

Python3中call方法介绍

切尔西携手YouTube 英超第一家共享视频球队诞生

商汤联手华科：提出文字检测模型GNNets，新颖模块可解决几何分布难题

(链表)反转链表Reverse List

很长时间没有来了

十年磨一剑，可重构计算架构将引领未来芯片市场

PyTorch中nn.Module类中call方法介绍

压缩和归档及vi的使用