模块

PyTorch 使用模块来表示神经网络。模块是：

• 有状态计算的基本构建块。 PyTorch 提供了一个强大的模块库，并使得定义新的自定义模块变得简单，从而可以轻松构建复杂的多层神经网络。

• 与PyTorch的 autograd 系统紧密集成。 模块使得为PyTorch的优化器指定可学习的参数变得简单，以便进行更新。

• 易于使用和转换。 模块易于保存和恢复，在CPU/GPU/TPU设备之间传输，修剪，量化等。

本笔记描述了模块，适用于所有 PyTorch 用户。由于模块在 PyTorch 中非常基础，本笔记中的许多主题在其他笔记或教程中有详细阐述，这里也提供了许多相关文档的链接。

一个简单的自定义模块

要开始，让我们来看一个更简单、自定义版本的 PyTorch 的 Linear 模块。 这个模块将其输入应用一个仿射变换。

import torch
from torch import nn

class MyLinear(nn.Module):
  def __init__(self, in_features, out_features):
    super().__init__()
    self.weight = nn.Parameter(torch.randn(in_features, out_features))
    self.bias = nn.Parameter(torch.randn(out_features))

  def forward(self, input):
    return (input @ self.weight) + self.bias

这个简单的模块具有以下模块的基本特征：

• 它继承自基础的 Module 类。 所有模块都应该子类化 Module 以便与其他模块组合使用。

• 它定义了一些用于计算的“状态”。 这里，状态由随机初始化的权重和偏置张量组成，这些张量定义了仿射变换。因为每个张量都被定义为参数，它们会被注册到模块中，并且会自动被跟踪并在调用参数()时返回。参数可以被认为是模块计算中“可学习”的方面（稍后会详细介绍）。请注意，模块不一定需要有状态，也可以是无状态的。

• 它定义了一个执行计算的forward()函数。 对于这个仿射变换模块，输入与weight参数（使用@简写符号）进行矩阵乘法，并加上bias参数以生成输出。更一般地，模块的forward()实现可以执行任意涉及任意数量输入和输出的计算。

这个简单的模块展示了如何将状态和计算打包在一起。可以构造并调用这个模块的实例：

m = MyLinear(4, 3)
sample_input = torch.randn(4)
m(sample_input)
: tensor([-0.3037, -1.0413, -4.2057], grad_fn=<AddBackward0>)

请注意，模块本身是可调用的，调用它将调用其 forward() 函数。 此名称指的是“前向传播”和“反向传播”的概念，这些概念适用于每个模块。 “前向传播”负责将模块表示的计算应用于给定的输入（如上例所示）。“反向传播”计算模块输出相对于其输入的梯度，这些梯度可以通过梯度下降方法用于“训练”参数。PyTorch 的 autograd 系统自动处理此反向传播计算，因此不需要为每个模块手动实现 backward() 函数。通过连续的前向/反向传播训练模块参数的过程在 使用模块进行神经网络训练 中有详细介绍。

可以通过调用 parameters() 或 named_parameters() 来迭代模块注册的完整参数集， 其中后者包括每个参数的名称：

for parameter in m.named_parameters():
  print(parameter)
: ('weight', Parameter containing:
tensor([[ 1.0597,  1.1796,  0.8247],
        [-0.5080, -1.2635, -1.1045],
        [ 0.0593,  0.2469, -1.4299],
        [-0.4926, -0.5457,  0.4793]], requires_grad=True))
('bias', Parameter containing:
tensor([ 0.3634,  0.2015, -0.8525], requires_grad=True))

一般来说，模块注册的参数是模块计算中应该被“学习”的方面。本笔记的后面部分展示了如何使用PyTorch的优化器之一来更新这些参数。然而，在此之前，让我们首先看看模块如何相互组合。

模块作为构建块

模块可以包含其他模块，这使得它们成为开发更复杂功能的有效构建块。 最简单的方法是使用Sequential模块。它允许我们将多个模块串联在一起：

net = nn.Sequential(
  MyLinear(4, 3),
  nn.ReLU(),
  MyLinear(3, 1)
)

sample_input = torch.randn(4)
net(sample_input)
: tensor([-0.6749], grad_fn=<AddBackward0>)

请注意，Sequential 会自动将第一个 MyLinear 模块的输出作为输入传递给 ReLU，并将该输出作为输入传递给第二个 MyLinear 模块。如所示，它仅限于具有单个输入和输出的模块的顺序链接。

通常建议为任何超出最简单用例的情况定义一个自定义模块，因为这提供了在模块计算中如何使用子模块的完全灵活性。

例如，这里是一个作为自定义模块实现的简单神经网络：

import torch.nn.functional as F

class Net(nn.Module):
  def __init__(self):
    super().__init__()
    self.l0 = MyLinear(4, 3)
    self.l1 = MyLinear(3, 1)
  def forward(self, x):
    x = self.l0(x)
    x = F.relu(x)
    x = self.l1(x)
    return x

该模块由两个“子模块”或“子模块”（l0 和 l1）组成，它们定义了神经网络的层，并在模块的 forward() 方法中用于计算。可以通过调用 children() 或 named_children() 来迭代模块的直接子模块：

net = Net()
for child in net.named_children():
  print(child)
: ('l0', MyLinear())
('l1', MyLinear())

要深入到不仅仅是直接子模块，modules() 和 named_modules() 递归遍历一个模块及其子模块：

class BigNet(nn.Module):
  def __init__(self):
    super().__init__()
    self.l1 = MyLinear(5, 4)
    self.net = Net()
  def forward(self, x):
    return self.net(self.l1(x))

big_net = BigNet()
for module in big_net.named_modules():
  print(module)
: ('', BigNet(
  (l1): MyLinear()
  (net): Net(
    (l0): MyLinear()
    (l1): MyLinear()
  )
))
('l1', MyLinear())
('net', Net(
  (l0): MyLinear()
  (l1): MyLinear()
))
('net.l0', MyLinear())
('net.l1', MyLinear())

有时，模块需要动态定义子模块。 ModuleList 和 ModuleDict 模块在这里很有用；它们 从列表或字典中注册子模块：

class DynamicNet(nn.Module):
  def __init__(self, num_layers):
    super().__init__()
    self.linears = nn.ModuleList(
      [MyLinear(4, 4) for _ in range(num_layers)])
    self.activations = nn.ModuleDict({
      'relu': nn.ReLU(),
      'lrelu': nn.LeakyReLU()
    })
    self.final = MyLinear(4, 1)
  def forward(self, x, act):
    for linear in self.linears:
      x = linear(x)
    x = self.activations[act](x)
    x = self.final(x)
    return x

dynamic_net = DynamicNet(3)
sample_input = torch.randn(4)
output = dynamic_net(sample_input, 'relu')

对于任何给定的模块，其参数包括其直接参数以及所有子模块的参数。 这意味着对 parameters() 和 named_parameters() 的调用将 递归地包含子参数，从而方便地优化网络中的所有参数：

for parameter in dynamic_net.named_parameters():
  print(parameter)
: ('linears.0.weight', Parameter containing:
tensor([[-1.2051,  0.7601,  1.1065,  0.1963],
        [ 3.0592,  0.4354,  1.6598,  0.9828],
        [-0.4446,  0.4628,  0.8774,  1.6848],
        [-0.1222,  1.5458,  1.1729,  1.4647]], requires_grad=True))
('linears.0.bias', Parameter containing:
tensor([ 1.5310,  1.0609, -2.0940,  1.1266], requires_grad=True))
('linears.1.weight', Parameter containing:
tensor([[ 2.1113, -0.0623, -1.0806,  0.3508],
        [-0.0550,  1.5317,  1.1064, -0.5562],
        [-0.4028, -0.6942,  1.5793, -1.0140],
        [-0.0329,  0.1160, -1.7183, -1.0434]], requires_grad=True))
('linears.1.bias', Parameter containing:
tensor([ 0.0361, -0.9768, -0.3889,  1.1613], requires_grad=True))
('linears.2.weight', Parameter containing:
tensor([[-2.6340, -0.3887, -0.9979,  0.0767],
        [-0.3526,  0.8756, -1.5847, -0.6016],
        [-0.3269, -0.1608,  0.2897, -2.0829, -2.0829],
        [ 2.6338,  0.9239,  0.6943, -1.5034]], requires_grad=True))
('linears.2.bias', Parameter containing:
tensor([ 1.0268,  0.4489, -0.9403,  0.1571], requires_grad=True))
('final.weight', Parameter containing:
tensor([[ 0.2509], [-0.5052], [ 0.3088], [-1.4951]], requires_grad=True))
('final.bias', Parameter containing:
tensor([0.3381], requires_grad=True))

也可以轻松地将所有参数移动到不同的设备或使用to()更改它们的精度：

# 将所有参数移动到CUDA设备
dynamic_net.to(device='cuda')

# 更改所有参数的精度
dynamic_net.to(dtype=torch.float64)

dynamic_net(torch.randn(5, device='cuda', dtype=torch.float64))
: tensor([6.5166], device='cuda:0', dtype=torch.float64, grad_fn=<AddBackward0>)

更一般地，可以通过使用apply()函数将任意函数递归地应用于模块及其子模块。例如，要将自定义初始化应用于模块及其子模块的参数：

# 定义一个函数来初始化线性权重。
# 注意，这里使用no_grad()来避免在autograd图中跟踪此计算。
@torch.no_grad()
def init_weights(m):
  if isinstance(m, nn.Linear):
    nn.init.xavier_normal_(m.weight)
    m.bias.fill_(0.0)

# 递归地在模块及其子模块上应用该函数。
dynamic_net.apply(init_weights)

这些示例展示了如何通过模块组合形成复杂的神经网络并方便地进行操作。为了允许快速且轻松地构建神经网络，PyTorch 在 torch.nn 命名空间中提供了一个大型的高性能模块库，这些模块执行常见的神经网络操作，如池化、卷积、损失函数等。

在下一节中，我们将给出一个完整的训练神经网络的示例。

更多信息，请查看：

• PyTorch 提供的模块库: torch.nn

• 定义神经网络模块: https://pytorch.org/tutorials/beginner/examples_nn/polynomial_module.html

使用模块进行神经网络训练

一旦网络构建完成，它就需要进行训练，并且可以使用 PyTorch 的其中一个优化器轻松优化其参数，例如来自 torch.optim：

# 创建网络（来自上一节）和优化器
net = Net()
optimizer = torch.optim.SGD(net.parameters(), lr=1e-4, weight_decay=1e-2, momentum=0.9)

# 运行一个示例训练循环，“教”网络
# 输出常数零函数
for _ in range(10000):
  input = torch.randn(4)
  output = net(input)
  loss = torch.abs(output)
  net.zero_grad()
  loss.backward()
  optimizer.step()

# 训练后，将模块切换到评估模式以进行推理、计算性能指标等。
# （有关训练和评估模式的描述，请参见下文讨论）
...
net.eval()
...

在这个简化的示例中，网络学习简单地输出零，因为任何非零输出都会根据其绝对值被“惩罚”，通过使用 torch.abs() 作为损失函数。虽然这不是一个非常有趣的任务，但训练的关键部分都存在：

• 创建了一个网络。

• 创建一个优化器（在本例中，是一个随机梯度下降优化器），并将网络的参数与其关联。

• A training loop…

  • 获取一个输入，

  • 运行网络，

  • 计算损失，

  • 将网络参数的梯度归零，

  • 调用 loss.backward() 来更新参数的梯度，

  • 调用 optimizer.step() 将梯度应用于参数。

在上述代码片段运行后，请注意网络的参数已经发生了变化。特别是，检查l1的weight参数的值，可以看到其值现在更接近于0（正如预期的那样）：

print(net.l1.weight)
: Parameter containing:
tensor([[-0.0013],
        [ 0.0030],
        [-0.0008]], requires_grad=True)

请注意，上述过程完全是在网络模块处于“训练模式”时完成的。模块默认处于训练模式，并且可以使用train()和 eval()在训练和评估模式之间切换。它们的行为可能会根据所处的模式而有所不同。例如，BatchNorm模块在训练期间维护一个运行均值和方差，当模块处于评估模式时不会更新这些值。通常，模块在训练期间应处于训练模式，并且仅在推理或评估时切换到评估模式。下面是一个自定义模块的示例，它在两种模式下的行为不同：

class ModalModule(nn.Module):
  def __init__(self):
    super().__init__()

  def forward(self, x):
    if self.training:
      # 仅在训练模式下添加一个常数。
      return x + 1.
    else:
      return x


m = ModalModule()
x = torch.randn(4)

print('训练模式输出: {}'.format(m(x)))
: tensor([1.6614, 1.2669, 1.0617, 1.6213, 0.5481])

m.eval()
print('评估模式输出: {}'.format(m(x)))
: tensor([ 0.6614,  0.2669,  0.0617,  0.6213, -0.4519])

训练神经网络通常可能会很棘手。更多信息，请查看：

• 使用优化器: https://pytorch.org/tutorials/beginner/examples_nn/two_layer_net_optim.html.

• 神经网络训练: https://pytorch.org/tutorials/beginner/blitz/neural_networks_tutorial.html

• 自动求导简介: https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html

模块状态

在上一节中，我们演示了训练模块的“参数”，即计算的可学习方面。 现在，如果我们想将训练好的模型保存到磁盘，可以通过保存其state_dict（即“状态字典”）来实现：

# 保存模块
torch.save(net.state_dict(), 'net.pt')

...

# 稍后加载模块
new_net = Net()
new_net.load_state_dict(torch.load('net.pt'))
: <所有键匹配成功>

模块的 state_dict 包含影响其计算的状态。这包括但不限于模块的参数。对于某些模块，除了参数之外，可能还需要有影响模块计算但不可学习的状态。对于这种情况，PyTorch 提供了“缓冲区”的概念，包括“持久性”和“非持久性”缓冲区。以下是模块可以拥有的各种类型状态的概述：

• 参数：计算的可学习方面；包含在 state_dict 中

• 缓冲区：计算中不可学习的方面

  • 持久化缓冲区：包含在state_dict中（即在保存和加载时序列化）

  • 非持久性缓冲区：不包含在state_dict中（即不进行序列化）

作为一个使用缓冲区的激励示例，考虑一个维护运行平均值的简单模块。我们希望当前的运行平均值被视为模块的state_dict的一部分，以便在加载模块的序列化形式时恢复它，但我们不希望它是可学习的。 这个代码片段展示了如何使用register_buffer()来实现这一点：

class RunningMean(nn.Module):
  def __init__(self, num_features, momentum=0.9):
    super().__init__()
    self.momentum = momentum
    self.register_buffer('mean', torch.zeros(num_features))
  def forward(self, x):
    self.mean = self.momentum * self.mean + (1.0 - self.momentum) * x
    return self.mean

现在，运行均值的当前值被视为模块的state_dict的一部分，并且在从磁盘加载模块时将正确恢复：

m = RunningMean(4)
for _ in range(10):
  input = torch.randn(4)
  m(input)

print(m.state_dict())
: OrderedDict([('mean', tensor([ 0.1041, -0.1113, -0.0647,  0.1515]))]))

# 序列化形式将包含'mean'张量
torch.save(m.state_dict(), 'mean.pt')

m_loaded = RunningMean(4)
m_loaded.load_state_dict(torch.load('mean.pt'))
assert(torch.all(m.mean == m_loaded.mean))

如前所述，可以通过将缓冲区标记为非持久性来将其排除在模块的 state_dict 之外：

self.register_buffer('unserialized_thing', torch.randn(5), persistent=False)

持久性和非持久性缓冲区都会受到使用to()应用的模型范围的设备/dtype更改的影响：

# 将所有模块参数和缓冲区移动到指定的设备/数据类型
m.to(device='cuda', dtype=torch.float64)

可以使用 buffers() 或 named_buffers() 来迭代模块的缓冲区。

for buffer in m.named_buffers():
  print(buffer)

以下类展示了在模块中注册参数和缓冲区的各种方法：

class StatefulModule(nn.Module):
  def __init__(self):
    super().__init__()
    # 将 nn.Parameter 设置为模块的属性会自动将张量注册为模块的参数。
    # 作为模块的参数。
    self.param1 = nn.Parameter(torch.randn(2))

    # 注册参数的基于字符串的替代方法。
    self.register_parameter('param2', nn.Parameter(torch.randn(3)))

    # 保留 "param3" 属性作为参数，防止它被设置为
    # 除了参数之外的任何东西。像这样的 "None" 条目将不会出现在模块的 state_dict 中。
    self.register_parameter('param3', None)

    # 注册参数列表。
    self.param_list = nn.ParameterList([nn.Parameter(torch.randn(2)) for i in range(3)])

    # 注册参数字典。
    self.param_dict = nn.ParameterDict({
      'foo': nn.Parameter(torch.randn(3)),
      'bar': nn.Parameter(torch.randn(4))
    })

    # 注册持久缓冲区（出现在模块的 state_dict 中）。
    self.register_buffer('buffer1', torch.randn(4), persistent=True)

    # 注册非持久缓冲区（不出现在模块的 state_dict 中）。
    self.register_buffer('buffer2', torch.randn(5), persistent=False)

    # 保留 "buffer3" 属性作为缓冲区，防止它被设置为
    # 除了缓冲区之外的任何东西。像这样的 "None" 条目将不会出现在模块的 state_dict 中。
    self.register_buffer('buffer3', None)

    # 添加子模块会将其参数注册为模块的参数。
    self.linear = nn.Linear(2, 3)

m = StatefulModule()

# 保存和加载 state_dict。
torch.save(m.state_dict(), 'state.pt')
m_loaded = StatefulModule()
m_loaded.load_state_dict(torch.load('state.pt'))

# 请注意，非持久缓冲区 "buffer2" 和保留属性 "param3" 和 "buffer3" 不会
# 出现在 state_dict 中。
print(m_loaded.state_dict())
: OrderedDict([('param1', tensor([-0.0322,  0.9066])),
               ('param2', tensor([-0.4472,  0.1409,  0.4852])),
               ('buffer1', tensor([ 0.6949, -0.1944,  1.2911, -2.1044])),
               ('param_list.0', tensor([ 0.4202, -0.1953])),
               ('param_list.1', tensor([ 1.5299, -0.8747])),
               ('param_list.2', tensor([-1.6289,  1.4898])),
               ('param_dict.bar', tensor([-0.6434,  1.5187,  0.0346, -0.4077])),
               ('param_dict.foo', tensor([-0.0845, -1.4324,  0.7022])),
               ('linear.weight', tensor([[-0.3915, -0.6176],
                                         [ 0.6062, -0.5992],
                                         [ 0.4452, -0.2843]])),
               ('linear.bias', tensor([-0.3710, -0.0795, -0.3947]))])

更多信息，请查看：

• 保存和加载: https://pytorch.org/tutorials/beginner/saving_loading_models.html

• 序列化语义: https://pytorch.org/docs/main/notes/serialization.html

• 什么是状态字典？ https://pytorch.org/tutorials/recipes/recipes/what_is_state_dict.html

模块初始化

默认情况下，torch.nn 提供的模块的参数和浮点缓冲区在模块实例化期间使用32位浮点值在CPU上进行初始化，采用一种历史上表现良好的初始化方案。对于某些用例，可能希望使用不同的数据类型（例如，GPU）或初始化技术进行初始化。

示例：

# 直接在GPU上初始化模块。
m = nn.Linear(5, 3, device='cuda')

# 使用16位浮点参数初始化模块。
m = nn.Linear(5, 3, dtype=torch.half)

# 跳过默认参数初始化并执行自定义（例如正交）初始化。
m = torch.nn.utils.skip_init(nn.Linear, 5, 3)
nn.init.orthogonal_(m.weight)

请注意，上述演示的设备和dtype选项也适用于为模块注册的任何浮点缓冲区：

m = nn.BatchNorm2d(3, dtype=torch.half)
print(m.running_mean)
: tensor([0., 0., 0.], dtype=torch.float16)

虽然模块编写者可以使用任何设备或数据类型来初始化其自定义模块中的参数，但良好的做法是默认使用dtype=torch.float和device='cpu'。此外，您可以通过遵循上述所有torch.nn模块所遵循的约定，为您的自定义模块在这些方面提供完全的灵活性。

• 提供一个 device 构造函数关键字参数，该参数适用于模块注册的任何参数/缓冲区。

• 提供一个应用于模块注册的任何参数/浮点缓冲区的dtype构造函数关键字参数。

• 仅在模块的构造函数中使用初始化函数（即来自 torch.nn.init 的函数）对参数和缓冲区进行初始化。请注意，仅在使用 skip_init() 时才需要这样做；有关解释，请参见 此页面。

更多信息，请查看：

• 跳过模块参数初始化: https://pytorch.org/tutorials/prototype/skip_param_init.html

模块钩子

在使用模块进行神经网络训练中，我们演示了模块的训练过程，该过程迭代地执行前向和后向传递，并在每次迭代中更新模块参数。为了更好地控制这一过程，PyTorch 提供了“钩子”，可以在前向或后向传递期间执行任意计算，甚至可以根据需要修改传递方式。此功能的一些有用示例包括调试、可视化激活、深入检查梯度等。钩子可以添加到您自己未编写的模块中，这意味着此功能可以应用于第三方或 PyTorch 提供的模块。

PyTorch 提供了两种类型的模块钩子：

• 前向钩子在正向传递期间被调用。它们可以通过 register_forward_pre_hook() 和 register_forward_hook() 为给定模块安装。 这些钩子将分别在调用前向函数之前和之后被调用。 或者，这些钩子可以通过类似的 register_module_forward_pre_hook() 和 register_module_forward_hook() 函数全局安装到所有模块。

• 反向钩子在反向传播过程中被调用。它们可以通过 register_full_backward_pre_hook() 和 register_full_backward_hook() 安装。 这些钩子将在该模块的反向传播计算完成后被调用。 register_full_backward_pre_hook() 将允许用户访问输出梯度， 而 register_full_backward_hook() 将允许用户访问输入和输出的梯度。或者，它们可以通过 register_module_full_backward_hook() 和 register_module_full_backward_pre_hook() 全局安装到所有模块。

所有钩子都允许用户返回一个更新后的值，该值将在剩余的计算过程中使用。因此，这些钩子可以用于在常规模块的前向/后向过程中执行任意代码，或者在不改变模块的forward()函数的情况下修改某些输入/输出。

下面是一个演示使用前向和后向钩子的示例：

```html

torch.manual_seed(1)

def forward_pre_hook(m, inputs):
  # 允许在正向传递之前检查和修改输入。
  # 注意输入总是包装在一个元组中。
  input = inputs[0]
  return input + 1.

def forward_hook(m, inputs, output):
  # 允许在正向传递之后检查输入/输出并修改输出。
  # 注意输入总是包装在一个元组中，而输出则按原样传递。

  # 残差计算，类似于ResNet。
  return output + inputs[0]

def backward_hook(m, grad_inputs, grad_outputs):
  # 允许检查grad_inputs / grad_outputs并修改
  # 用于反向传递的grad_inputs。注意grad_inputs和
  # grad_outputs总是包装在元组中。
  new_grad_inputs = [torch.ones_like(gi) * 42. for gi in grad_inputs]
  return new_grad_inputs

# 创建示例模块和输入。
m = nn.Linear(3, 3)
x = torch.randn(2, 3, requires_grad=True)

# ==== 演示正向钩子。 ====
# 在添加钩子之前和之后通过模块运行输入。
print('output with no forward hooks: {}'.format(m(x)))
: output with no forward hooks: tensor([[-0.5059, -0.8158,  0.2390],
                                        [-0.0043,  0.4724, -0.1714]], grad_fn=<AddmmBackward>)

# 注意修改后的输入导致不同的输出。
forward_pre_hook_handle = m.register_forward_pre_hook(forward_pre_hook)
print('output with forward pre hook: {}'.format(m(x)))
: output with forward pre hook: tensor([[-0.5752, -0.7421,  0.4942],
                                        [-0.0736,  0.5461,  0.0838]], grad_fn=<AddmmBackward>)

# 注意修改后的输出。
forward_hook_handle = m.register_forward_hook(forward_hook)
print('output with both forward hooks: {}'.format(m(x)))
: output with both forward hooks: tensor([[-1.0980,  0.6396,  0.4666],
                                          [ 0.3634,  0.6538,  1.0256]], grad_fn=<AddBackward0>)

# 移除钩子；注意这里的输出与添加钩子之前的输出匹配。
forward_pre_hook_handle.remove()
forward_hook_handle.remove()
print('output after removing forward hooks: {}'.format(m(x)))
: output after removing forward hooks: tensor([[-0.5059, -0.8158,  0.2390],
                                               [-0.0043,  0.4724, -0.1714]], grad_fn=<AddmmBackward>)

# ==== 演示反向钩子。 ====
m(x).sum().backward()
print('x.grad with no backwards hook: {}'.format(x.grad))
: x.grad with no backwards hook: tensor([[ 0.4497, -0.5046,  0.3146],
                                         [ 0.4497, -0.5046,  0.3146]])

# 在再次运行反向传递之前清除梯度。
m.zero_grad()
x.grad.zero_()

m.register_full_backward_hook(backward_hook)
m(x).sum().backward()
print('x.grad with backwards hook: {}'.format(x.grad))
: x.grad with backwards hook: tensor([[42., 42., 42.],
                                      [42.,</span

高级功能

PyTorch 还提供了几个更高级的功能，这些功能旨在与模块一起使用。所有这些功能都适用于自定义编写的模块，但有一点需要注意的是，某些功能可能需要模块符合特定的约束条件才能得到支持。关于这些功能及其相应要求的深入讨论，可以在以下链接中找到。

分布式训练

PyTorch中存在多种分布式训练方法，既可以用于使用多个GPU进行扩展训练，也可以用于跨多台机器进行训练。查看分布式训练概述页面以获取有关如何利用这些方法的详细信息。

性能分析

PyTorch Profiler 可以用于识别模型中的性能瓶颈。它测量并输出内存使用和时间消耗的性能特征。

通过量化提高性能

将量化技术应用于模块可以通过使用比浮点精度更低的位宽来提高性能和内存使用率。查看PyTorch提供的各种量化机制，请点击这里。

通过剪枝改善内存使用

大型深度学习模型通常参数过多，导致内存使用量高。为了应对这一问题，PyTorch 提供了模型剪枝的机制，这可以在保持任务准确性的同时帮助减少内存使用。剪枝教程描述了如何利用 PyTorch 提供的剪枝技术，或在必要时定义自定义剪枝技术。

参数化

对于某些应用，在模型训练期间限制参数空间可能是有益的。例如，强制学习参数的正交性可以改善RNN的收敛性。PyTorch提供了一种机制，用于应用参数化，并进一步允许定义自定义约束。

使用FX转换模块

PyTorch 的 FX 组件提供了一种灵活的方式，通过直接操作模块计算图来转换模块。这可以用于以编程方式生成或操作模块，以满足广泛的用例需求。要探索 FX，请查看这些使用 FX 的示例，包括 卷积 + 批量归一化融合 和 CPU 性能分析。