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

文章目录

• 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_c和linear。net_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()。