【torch.nn.Module】神经网络模块类及其方法

torch.nn.Module

类的定义

torch.nn.Module 是所有神经网络模块的基类。自定义的模型也应该子类化这个类。该模块还可以包含其他模块,允许将它们嵌套在树结构中。你可以为子模块指定常规属性:

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

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
        x = F.relu(self.conv1(x))
        return F.relu(self.conv2(x))

以这种方式分配的子模块将被注册,并且在调用to()时也将转换它们的参数,等等。根据上面的例子,在赋值给子类之前,必须对父类进行__init__()调用。

  • 参数
    • training (bool): 布尔值表示该模块是处于训练模式还是评估模式。

相关方法

add_module

add_module(name, module)
  • 功能
    • 将子模块添加到当前模块。可以使用给定的名称将模块作为属性访问。
  • 参数
    • name (str) 子模块的名称。可以使用给定的名称从该模块访问子模块。
    • module (module) 要添加到模块中的子模块。

apply

apply(fn)
  • 功能: 将fn递归应用于每个子模块(由.children()返回)和self。典型的用法包括初始化模型的参数(参见torch.nn.init)。
  • 参数
    • fn (Module -> None) 要应用到每个子模块的函数
  • 返回
    • self (Module)
  • 例子
>>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[1., 1.],
        [1., 1.]], requires_grad=True)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)

bfloat

bfloat16()
  • 功能

    将所有浮点参数和缓冲区转换为bfloat16数据类型。

  • 返回值

    self(Module)

buffers

buffers(recurse=True)
  • 功能

    返回模块缓冲区上的迭代器。

  • 参数

    • recurse (bool) 如果为True,则产生该模块和所有子模块的缓冲区。否则,只生成该模块的直接成员的缓冲区。
  • Yield

    • torch.Tensor 模型缓冲器
  • 例子

for buf in model.buffers():
    print(type(buf), buf.size())

children

children()
  • 功能

    返回直接子模块上的迭代器。

  • Yield

    • Module - 子模型
  • 返回

    • iterator[Module]

cpu

cpu()
  • 功能:将所有模型参数和缓冲区移动到CPU。
  • 返回值: self(Module)

cuda

cuda(device=None)

将所有模型参数和缓冲区移动到GPU。这也使得关联的参数和缓冲区是不同的对象。因此,如果模块在优化时将驻留在GPU上,则应该在构造优化器之前调用它。

  • 参数
    • device (int, optional) 如果指定,所有参数将被复制到该设备
  • 返回值:self(Module)

double

double()
  • 功能:将所有浮点参数和缓冲区转换为double类型。
  • 返回值:self(Module)

eval

eval()
  • 功能 :

    将模块设置为评估模式。这只对某些模块有影响。模块在训练/评估模式下的行为,如Dropout, BatchNorm等,请参见具体模块的文档。这相当于self.train(False)

  • 返回 : self(Moudle)

extra_repr

extra_repr()
  • 功能

    设置模块的额外表示形式要打印定制的额外信息,您应该在自己的模块中重新实现此方法。单行和多行字符串都可以接受。

  • 返回: str

float

float()
  • 功能: ****将所有浮点参数和缓冲区转换为浮点数据类型。
  • 返回:self(Moudle)

forward

forward(*input)
  • 功能:定义每次调用时执行的计算。应该被所有子类重写。

get_buffer

get_buffer(target)
  • 功能: 如果存在,则返回目标给定的缓冲区,否则抛出错误。有关此方法的功能以及如何正确指定target的更详细说明,请参阅get_submodule的文档字符串。
  • 参数
    • target(str) 要查找的缓冲区的完全限定字符串名称。
  • 返回: target引用的缓冲区
  • 返回类型: torch.Tensor

get_extra_state

get_extra_state()
  • 功能

    返回要包含在模块的state_dict中的任何额外状态。如果您需要存储额外的状态,则为您的模块实现这个和相应的set_extra_state()。在构建模块的state_dict()时调用此函数。

    注意,额外的状态应该是可pickle的,以确保state_dict的序列化工作。我们仅为序列化张量提供向后兼容性保证;如果其他对象的序列化pickle形式发生变化,则可能会破坏向后兼容性。

  • 返回: 任何要存储在模块的state_dict中的额外状态

  • 返回类型:object

get_parameter

get_parameter(target)
  • 功能:如果存在,则返回target给定的参数,否则抛出错误。有关此方法的功能以及如何正确指定target的更详细说明,请参阅get_submodule的文档字符串。
  • 参数:
    • target (str): 要查找的Parameter的字符串名称
  • 返回: target引用的参数
  • 返回类型:torch.nn.Parameter

get_submodule

get_submodule(target)
  • 功能

如果子模块存在,则返回target给出的子模块,否则抛出错误。

例如,假设你有一个nn。模块A是这样的:

A(
    (net_b): Module(
        (net_c): Module(
            (conv): Conv2d(16, 33, kernel_size=(3, 3), stride=(2, 2))
        )
        (linear): Linear(in_features=100, out_features=200, bias=True)
    )
)

(该图显示了一个**nn.Module A.** A有一个嵌套的子模块net_b,它本身有两个子模块net_clinearnet_c有一个子模块conv。)

get_submodule的运行时间受模块在target中的嵌套程度的限制。对named_modules的查询实现了相同的结果,但是它的可传递模块数量是O(N)。因此,为了简单地检查是否存在子模块,应该始终使用get_submodule

  • 参数
    • target (str): 要查找的子模型的字符串名称
  • 返回: target引用的子模块
  • 返回类型:torch.nn.module

half

half()
  • 功能:将所有浮点参数和缓冲区转换为half类型
  • 返回:self
  • 返回类型:Module

ipu

ipu(device=None)
  • 功能:

移动所有模型参数和缓冲区到IPU。这也使得关联的参数和缓冲区是不同的对象。因此,如果模块将在IPU上进行优化,则应该在构造优化器之前调用它。

  • 参数:
    • device (int, optional) : 如果指定,所有参数将被复制到该设备
  • 返回:self
  • 返回类型:Module

load_state_dict

load_state_dict(state_dict, strict=True)
  • 功能:将参数和缓冲区从state_dict复制到此模块及其子模块中。如果strict为True,则state_dict的键必须与该模块的state_dict()函数返回的键完全匹配。
  • 参数
    • state_dict (dict)- 包含参数和持久缓冲区的字典
    • strict (bool, optional)-**是否严格强制state_dict中的键与本模块的state_dict()函数返回的键匹配。默认值: true
  • 返回:
    • Missing_keys是一个包含缺失键的STR列表
    • unexpected ted_keys是一个包含意外键的STR列表
  • 返回类型:带有missing_keys和expected_keys字段的命名元组

modules

modules()
  • 功能:返回网络中所有模块的迭代器。
  • Yield: Module –网络中的模型
  • 返回类型:iterator[Module]
  • 例子
>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
...     print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)

named_buffers

named_buffers(prefix='', recurse=True)
  • 功能:返回模块缓冲区上的迭代器,产生缓冲区的名称和缓冲区本身**。**
  • 参数
    • prefix (str) : 所有缓冲区名称前的前缀。
    • recurse (bool) : 如果为True,则产生该模块和所有子模块的缓冲区。否则,只生成该模块的直接成员的缓冲区。
  • Yield:
    • (str, torch.Tensor) 包含名称和缓冲区的元组
  • 返回类型:iterator[Tuple[str,Tensor]]
  • 例子
>>> for name, buf in self.named_buffers():
>>>    if name in ['running_var']:
>>>        print(buf.size())

named_children

named_children()
  • 功能:返回直接子模块的迭代器,包括模块名和模块本身**。**
  • Yield:
    • (str,Module) 包含名称和子模型的元组
  • 返回类型:iterator[Tuple[str,Module]]
  • 例子
>>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)

named_modules

named_modules(memo=None, prefix='', remove_duplicate=True)
  • 功能:返回网络中所有模块的迭代器,包括模块名和模块本身。
  • 参数
    • memo(Optional[Set[Module]]) -一个备忘录,用来存储已经添加到结果中的模块集
    • prefix (str) -将添加到模块名称的前缀
    • remove_duplicate (bool)——是否在结果中删除重复的模块实例
  • Yield
    • (str, Module) – 名称和模块的元组
  • 例子
>>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
...     print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))

named_parameters

named_parameters(prefix='', recurse=True)
  • 功能:返回模块参数的迭代器,包括形参名称和形参本身。
  • 参数
    • Prefix (str) -所有参数名称前的前缀。
    • recurse (bool) -如果为True,则产生该模块及其所有子模块的参数。否则,只生成该模块的直接成员的参数。
  • Yield:
    • (str,Parameter) 包含名称和参数的元组
  • 返回类型:iterator[Tuple[str,Parameter]]
  • 例子
>>> for name, param in self.named_parameters():
>>>    if name in ['bias']:
>>>        print(param.size())

parameters

parameters(recurse=True)
  • 功能:返回模块参数的迭代器。这通常被传递给优化器。
  • 参数
    • recurse (bool) -如果为True,则产生该模块及其所有子模块的参数。否则,只生成该模块的直接成员的参数。
  • Yield:
    • Parameter- 模型参数
  • 返回类型:iterator[Parameter]
  • 例子
>>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)

register_backward_hook

register_backward_hook(hook)
  • 功能:在模型上注册一个hook。 该函数已弃用,取而代之的是register_full_backward_hook(),该函数的行为将在未来版本中改变。
  • 返回:一个句柄,可用于通过调用**handle.remove()**来删除添加的钩子
  • 返回类型:torch.utils.hooks.RemovableHandle

register_buffer

register_buffer(name, tensor, persistent=True)
  • 功能:

向模块中添加缓冲区。

这通常用于注册一个不应被视为模型参数的缓冲区。例如,BatchNorm的running_mean不是一个参数,而是模块状态的一部分。默认情况下,缓冲区是持久的,将与参数一起保存。这种行为可以通过将persistent设置为False来改变。持久缓冲区和非持久缓冲区之间的唯一区别是,后者不会成为该模块的state_dict的一部分。

缓冲区可以作为使用给定名称的属性访问。

  • 参数
    • name(str) : 缓冲区的名称。可以使用给定的名称从该模块访问缓冲区
    • tensor (Tensor or None) : 要注册的缓冲区。如果为None,则在缓冲区上运行的操作,例如cuda,将被忽略。如果为None,则该缓冲区不包含在模块的state_dict中。
    • persistent (bool) : 缓冲区是否是这个模块的state_dict的一部分。
  • 例子
self.register_buffer('running_mean', torch.zeros(num_features))

register_forward_hook

register_forward_hook(hook)
  • 功能:

在模块上注册一个前向钩子。

每当forward()计算出输出后,钩子都会被调用。文件须有下列signature:

hook(module, input, output) -> None or modified output

输入只包含给定给模块的位置参数。关键字参数不会传递给钩子,而只传递给forward。钩子可以修改输出。它可以就地修改输入,但不会对forward产生影响,因为它是在forward()被调用之后调用的。

  • 返回:一个句柄,可以通过调用handle.remove()来删除添加的钩子。
  • 返回类型:torch.utils.hooks.RemovableHandle

register_forward_pre_hook

register_forward_pre_hook(hook)
  • 功能:在模块上注册一个pre-hook。每次forward()调用之前都会调用钩子。文件须有下列签署:
hook(module, input) -> None or modified input

输入只包含给定给模块的位置参数。关键字参数不会传递给钩子,而只传递给forward。钩子可以修改输入。用户可以在钩子中返回一个元组或单个修改后的值。如果返回单个值(除非该值已经是元组),则将该值包装到一个元组中。

  • 返回:一个句柄,可以通过调用handle.remove()来删除添加的钩子。
  • 返回类型:torch.utils.hooks.RemovableHandle

register_full_backward_hook

register_full_backward_hook(hook)
  • 功能:在模块上注册一个后向钩子。每次计算模块输入的梯度时,钩子都会被调用。钩子应该有以下signature:

    hook(module, grad_input, grad_output) -> tuple(Tensor) or None
    
  • 返回:一个句柄,可以通过调用handle.remove()来删除添加的钩子。

  • 返回类型:torch.utils.hooks.RemovableHandle

register_load_state_dict_post_hook

register_load_state_dict_post_hook(hook)
  • 功能:注册一个post钩子,在模块的load_state_dict被调用后运行。
  • 返回:一个句柄,可以通过调用handle.remove()来删除添加的钩子。
  • 返回类型:torch.utils.hooks.RemovableHandle

register_module

register_module(name, module)
  • 功能:add_module()的别名

register_parameter

register_parameter(name, param)
  • 功能:向模块中添加参数。该参数可以作为使用给定名称的属性访问。
  • 参数
    • name (str) : 参数名称。可以使用给定的名称从该模块访问该参数
    • param(Parameter or None) : 要添加到模块中的参数。如果为None,那么运行在参数上的操作,比如cuda,将被忽略。如果为None,则该参数不包含在模块的state_dict中。

requires_grad_

requires_grad_(requires_grad=True)
  • 功能 :

更改autograd是否应该记录此模块中参数的操作。该方法就地设置参数的requires_grad属性。这种方法有助于冻结模块的部分以进行微调或单独训练模型的部分(例如,GAN训练)。关于.requires_grad_()和一些可能与之混淆的类似机制之间的比较,请参阅局部禁用梯度计算。

  • 参数
    • requires_grad (bool) autograd是否应该记录此模块中参数的操作。默认值: True。
  • 返回值
    • self
  • 返回类型
    • Moudule

set_extra_state

set_extra_state(state)
  • 功能: 该函数从load_state_dict()中调用,以处理在state_dict中发现的任何额外状态。如果您需要在state_dict中存储额外的状态,则为您的模块实现这个函数和相应的get_extra_state()
  • 参数
    • state (dict) – Extra state from the state_dict

share_memory

share_memory()
  • 功能:同**[torch.Tensor.share_memory_()](https://pytorch.org/docs/stable/generated/torch.Tensor.share_memory_.html#torch.Tensor.share_memory_)**
  • 返回类型:T

state_dict

state_dict(*, destination: T_destination, prefix: str = '', keep_vars: bool = False) → T_destination
  • 功能:返回一个字典,其中包含对模块整个状态的引用。参数和持久缓冲区(例如运行平均值)都包括在内。键是对应的参数和缓冲区名称。参数和缓冲区设置为None不包括在内。
  • 参数:
    • destination (dict, optional)** 如果提供,模块的状态将被更新到dict中,并返回相同的对象。否则,将创建并返回OrderedDict。默认值:None。
    • prefix (str, optional)** 添加到参数名和缓冲区名中的前缀,用于组成state_dict中的键。默认值:
    • keep_vars (bool, optional)**
  • 返回:包含模块的整个状态的字典
  • 返回类型:dict
  • 例子
>>> module.state_dict().keys()
['bias', 'weight']

to

to(device: Optional[Union[int, device]] = ..., dtype: Optional[Union[dtype, str]] = ..., non_blocking: bool = ...) → T
to(dtype: Union[dtype, str], non_blocking: bool = ...) → T
to(tensor: Tensor, non_blocking: bool = ...) → T
  • 功能:移动和/或强制转换参数和缓冲区。调用如下
to(device=None, dtype=None, non_blocking=False)
to(dtype, non_blocking=False)
to(tensor, non_blocking=False)
to(memory_format=torch.channels_last)

它的signature类似于torch.Tensor.to(),但只接受浮点或复杂dtypes。此外,此方法仅将浮点或复杂参数和缓冲区转换为dtype(如果给定)。如果给定,则整型参数和缓冲区将被移动到设备,但dtypes不变。当设置non_blocking时,如果可能的话,它会尝试相对于主机异步转换/移动,例如,将CPU张量与固定内存移动到CUDA设备。

  • 参数:

    • device (torch.device) 此模块中参数和缓冲区的所需设备
    • dtype (toch.dtype) 此模块中参数和缓冲区所需的浮点或复杂dtype
    • **tensor (torch.Tensor)**张量,其dtype和device是该模块中所有参数和缓冲区所需的dtype和device
    • memory_format (torch.memory_format) 该模块中4D参数和缓冲区所需的内存格式(只有关键字参数)
  • 返回:

    self(Moudule)

  • 例子

>>> linear = nn.Linear(2, 2)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]])
>>> linear.to(torch.double)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1913, -0.3420],
        [-0.5113, -0.2325]], dtype=torch.float64)
>>> gpu1 = torch.device("cuda:1")
>>> linear.to(gpu1, dtype=torch.half, non_blocking=True)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16, device='cuda:1')
>>> cpu = torch.device("cpu")
>>> linear.to(cpu)
Linear(in_features=2, out_features=2, bias=True)
>>> linear.weight
Parameter containing:
tensor([[ 0.1914, -0.3420],
        [-0.5112, -0.2324]], dtype=torch.float16)

>>> linear = nn.Linear(2, 2, bias=None).to(torch.cdouble)
>>> linear.weight
Parameter containing:
tensor([[ 0.3741+0.j,  0.2382+0.j],
        [ 0.5593+0.j, -0.4443+0.j]], dtype=torch.complex128)
>>> linear(torch.ones(3, 2, dtype=torch.cdouble))
tensor([[0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j],
        [0.6122+0.j, 0.1150+0.j]], dtype=torch.complex128)

to_empty

to_empty(*, device)
  • 功能:将参数和缓冲区移动到指定的设备,而不复制存储。
  • 参数:
    • device (torch.device) – 此模块中参数和缓冲区的所需设备。
  • 返回:self(Module)

train

train(mode=True)
  • 功能:将模块设置为train模式。这只对某些模块有影响。具体模块在训练/评估模式下的行为,如Dropout, BatchNorm等,请参见具体模块的文档。
  • 参数:
    • mode (bool) - 是否设置train模式(True)或evaluation模式(False)。默认值:True。
  • 返回:self(Module)

type

type(dst_type)
  • 功能
    • 将所有参数和缓冲区转换为dst_type。
  • 参数
    • dst_type (type or string) 期望的类型
  • 返回
    • self(Module)

xpu

xpu(device=None)
  • 功能

    将所有模型参数和缓冲区移动到XPU。这也使得关联的参数和缓冲区不同的对象。因此,如果模块在优化时将驻留在XPU上,则应该在构造优化器之前调用它。

  • 参数

    • device (int, optional)** 如果指定,所有参数将被复制到该设备
  • 返回

    • self(Module)

zero_grad

zero_grad(set_to_none=False)
  • 功能: ****将所有模型参数的梯度设置为零。有关更多上下文,请参阅torch. optimizer下的类似函数。
  • 参数
    • set_to_none (bool) 而不是设置为零,设置梯度为None。详情请参阅torch.optim.Optimizer.zero_grad()

猜你喜欢

转载自blog.csdn.net/zyw2002/article/details/128273928