模块

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

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

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

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

此说明介绍了模块，适用于所有 PyTorch 用户。由于模块是 PyTorch 的基础，因此 本说明中的许多主题在其他说明或教程中进行了详细说明，并链接到其中的许多文档 也在此处提供。

一个简单的自定义模块

首先，让我们看看 PyTorch 模块的更简单的自定义版本。 此模块对其 input 应用仿射变换。

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

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

• 它继承自 Base Module 类。所有模块都应该是子类，以便与其他模块进行可组合性。

• 它定义了一些用于计算的 “state”。在这里，状态由随机初始化的张量和定义仿射的张量组成 转型。由于它们中的每一个都定义为 ，因此它们都已注册到模块中，并且将自动跟踪并从调用中返回 到 .参数可以是 考虑了模块计算的 “可学习” 方面（稍后会详细介绍）。请注意，模块 不需要有 state，也可以是无状态的。weightbias

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

这个简单的模块演示了模块如何将 state 和 computation 打包在一起。此模块的实例可以是 constructed 并调用：

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

请注意，模块本身是可调用的，并且调用它会调用其函数。 此名称引用了适用于每个模块的 “forward pass” 和 “backward pass” 的概念。 “forward pass” 负责应用模块表示的计算 添加到给定的 Input（如上面的代码段所示）。“backward pass” 计算的梯度为 模块输出，可用于通过 Gradient “训练” 参数 descent 方法。PyTorch 的 autograd 系统会自动处理这种向后传递计算，因此它 不需要为每个模块手动实现一个函数。培训过程 通过连续向前/向后传递的模块参数在 使用模块进行神经网络训练 中详细介绍。forward()backward()

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

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 的优化器之一更新这些参数。 然而，在我们开始之前，让我们首先研究一下模块是如何相互组合的。

作为构建块的模块

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

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>)

请注意，会自动将第一个模块的输出作为 input 馈送 到 中，并将其输出作为输入到第二个模块中。如 所示，它仅限于具有单个输入和输出的模块的 In-Order Chaining。MyLinearMyLinear

一般来说，建议为最简单的用例之外的任何内容定义自定义模块，因为这提供了 在如何将子模块用于模块的计算方面具有完全的灵活性。

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

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

此模块由两个“子”或“子模块”（ 和 ）组成，它们定义了 神经网络 和 Neural Network 的 API 中，用于模块方法中的计算。立即的 模块的子项可以通过调用 or 进行迭代：l0l1forward()

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

要更深入地了解直接的子模块，并递归地迭代模块及其子模块：

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())

有时，模块需要动态定义子模块。 和 modules 在这里很有用;他们 从 list 或 dict 中注册 submodules：

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')

对于任何给定的模块，它的参数由它的直接参数以及所有子模块的参数组成。 这意味着对 和 将 递归包含子参数，从而可以方便地优化网络内的所有参数：

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.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))

将所有参数移动到不同的设备或更改其精度也很容易：

# Move all parameters to a CUDA device
dynamic_net.to(device='cuda')

# Change precision of all parameters
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>)

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

# Define a function to initialize Linear weights.
# Note that no_grad() is used here to avoid tracking this computation in the autograd graph.
@torch.no_grad()
def init_weights(m):
  if isinstance(m, nn.Linear):
    nn.init.xavier_normal_(m.weight)
    m.bias.fill_(0.0)

# Apply the function recursively on the module and its submodules.
dynamic_net.apply(init_weights)

这些示例展示了如何通过模块组合方便地形成复杂的神经网络 数据处理。为了能够使用最少的样板代码快速轻松地构建神经网络，PyTorch 在命名空间中提供了一个大型高性能模块库，这些模块执行常见的神经 网络操作，如池化、卷积、损失函数等。

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

有关更多信息，请查看：

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

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

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

一旦构建了网络，就必须对其进行训练，并且其参数可以使用 PyTorch 的 优化器来自 ：

# Create the network (from previous section) and optimizer
net = Net()
optimizer = torch.optim.SGD(net.parameters(), lr=1e-4, weight_decay=1e-2, momentum=0.9)

# Run a sample training loop that "teaches" the network
# to output the constant zero function
for _ in range(10000):
  input = torch.randn(4)
  output = net(input)
  loss = torch.abs(output)
  net.zero_grad()
  loss.backward()
  optimizer.step()

# After training, switch the module to eval mode to do inference, compute performance metrics, etc.
# (see discussion below for a description of training and evaluation modes)
...
net.eval()
...

在这个简化的例子中，网络学习简单地输出零，因为任何非零输出都会根据 到它的绝对值。虽然这不是一项非常有趣的任务， 培训的关键部分包括：

• 将创建一个网络。

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

• 一个训练循环...

  • 获取输入，

  • 运行网络，

  • 计算损失，

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

  • 调用 loss.backward（） 来更新参数的梯度，

  • 调用 optimizer.step（） 将梯度应用于参数。

运行上述代码段后，请注意，网络的参数已更改。具体而言，检查 的 value of 的参数显示其值现在更接近 0（正如预期的那样）：l1weight

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

请注意，上述过程完全是在网络模块处于 “training mode” 时完成的。模块默认为 training 模式，并且可以使用 和 在训练模式和评估模式之间切换。它们的行为可能会有所不同，具体取决于它们所处的模式。例如，该模块在训练期间维护未更新的运行均值和方差 当模块处于评估模式时。通常，模块在训练期间应处于训练模式 并且仅切换到评估模式进行推理或评估。下面是自定义模块的示例 在两种模式之间表现不同：BatchNorm

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

  def forward(self, x):
    if self.training:
      # Add a constant only in training mode.
      return x + 1.
    else:
      return x


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

print('training mode output: {}'.format(m(x)))
: tensor([1.6614, 1.2669, 1.0617, 1.6213, 0.5481])

m.eval()
print('evaluation mode output: {}'.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

• autograd 简介：https://pytorch.org/tutorials/beginner/blitz/autograd_tutorial.html

模块状态

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

# Save the module
torch.save(net.state_dict(), 'net.pt')

...

# Load the module later on
new_net = Net()
new_net.load_state_dict(torch.load('net.pt'))
: <All keys matched successfully>

模块的 contains 状态会影响其计算。这包括但不限于 模块的参数。对于某些模块，在影响模块的参数之外具有 state 可能很有用 计算，但不可学习。对于这种情况，PyTorch 提供了“缓冲区”的概念，两者都是“持久”的 和 “non-persistent” 的 intent 的以下是模块可以具有的各种类型的状态的概述：state_dict

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

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

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

  • 非持久缓冲区：不包含在 （即被排除在序列化之外）state_dict

作为使用 buffers 的一个激励性示例，请考虑一个保持 running mean 的简单模块。我们希望 运行均值的当前值被视为模块的一部分，因此它将是 restored 的，但我们不希望它是可学习的。 此代码段演示如何使用来实现此目的：state_dict

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]))]))

# Serialized form will contain the 'mean' tensor
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)

持久缓冲区和非持久缓冲区都受到模型范围的设备/dtype 更改的影响，这些更改应用了：

# Moves all module parameters and buffers to the specified device / dtype
m.to(device='cuda', dtype=torch.float64)

可以使用 或 迭代模块的缓冲区。

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

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

class StatefulModule(nn.Module):
  def __init__(self):
    super().__init__()
    # Setting a nn.Parameter as an attribute of the module automatically registers the tensor
    # as a parameter of the module.
    self.param1 = nn.Parameter(torch.randn(2))

    # Alternative string-based way to register a parameter.
    self.register_parameter('param2', nn.Parameter(torch.randn(3)))

    # Reserves the "param3" attribute as a parameter, preventing it from being set to anything
    # except a parameter. "None" entries like this will not be present in the module's state_dict.
    self.register_parameter('param3', None)

    # Registers a list of parameters.
    self.param_list = nn.ParameterList([nn.Parameter(torch.randn(2)) for i in range(3)])

    # Registers a dictionary of parameters.
    self.param_dict = nn.ParameterDict({
      'foo': nn.Parameter(torch.randn(3)),
      'bar': nn.Parameter(torch.randn(4))
    })

    # Registers a persistent buffer (one that appears in the module's state_dict).
    self.register_buffer('buffer1', torch.randn(4), persistent=True)

    # Registers a non-persistent buffer (one that does not appear in the module's state_dict).
    self.register_buffer('buffer2', torch.randn(5), persistent=False)

    # Reserves the "buffer3" attribute as a buffer, preventing it from being set to anything
    # except a buffer. "None" entries like this will not be present in the module's state_dict.
    self.register_buffer('buffer3', None)

    # Adding a submodule registers its parameters as parameters of the module.
    self.linear = nn.Linear(2, 3)

m = StatefulModule()

# Save and load state_dict.
torch.save(m.state_dict(), 'state.pt')
m_loaded = StatefulModule()
m_loaded.load_state_dict(torch.load('state.pt'))

# Note that non-persistent buffer "buffer2" and reserved attributes "param3" and "buffer3" do
# not appear in the 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

模块初始化

默认情况下，由 提供的模块的参数和浮点缓冲区在 模块实例化为 CPU 上的 32 位浮点值，使用确定为 模块类型的历史表现良好。对于某些用例，可能需要使用不同的 dtype、设备（例如 GPU）或初始化技术。

例子：

# Initialize module directly onto GPU.
m = nn.Linear(5, 3, device='cuda')

# Initialize module with 16-bit floating point parameters.
m = nn.Linear(5, 3, dtype=torch.half)

# Skip default parameter initialization and perform custom (e.g. orthogonal) initialization.
m = torch.nn.utils.skip_init(nn.Linear, 5, 3)
nn.init.orthogonal_(m.weight)

请注意，上面演示的 device 和 dtype 选项也适用于任何已注册的浮点缓冲区 对于该模块：

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

虽然模块编写者可以使用任何设备或 dtype 来初始化其自定义模块中的参数，但好的做法是 以使用 和 默认情况下。或者，您可以提供充分的灵活性 在这些区域中，通过遵守上面演示的约定，所有模块都遵循：dtype=torch.floatdevice='cpu'

• 提供适用于模块注册的任何参数/缓冲区的构造函数 kwarg。device

• 提供一个构造函数 kwarg，该构造函数适用于由 模块。dtype

• 仅在 module 的构造函数。请注意，这仅在使用 时需要 ;有关说明，请参阅此页面。

有关更多信息，请查看：

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

模块钩子

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

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

• 在前向传递期间调用前向钩子。可以使用 和 为给定模块安装它们。 这些钩子将在调用 forward 函数之前和调用之后分别被调用。 或者，这些钩子可以全局安装用于所有具有 analogous 和 functions 的模块。

• 在向后传递期间调用向后钩子。它们可以使用 和 一起安装。 当计算出此 Module 的 backward 时，将调用这些钩子。将允许用户访问输出的梯度 while 将允许用户访问渐变 输入和输出。或者，它们可以全局安装用于所有带有 和 的模块。

所有钩子都允许用户返回一个更新的值，该值将在剩余的计算中使用。 因此，这些钩子可用于沿常规模块向前/向后执行任意代码，或者 修改一些 inputs/outputs，而不必更改模块的功能。forward()

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

torch.manual_seed(1)

def forward_pre_hook(m, inputs):
  # Allows for examination and modification of the input before the forward pass.
  # Note that inputs are always wrapped in a tuple.
  input = inputs[0]
  return input + 1.

def forward_hook(m, inputs, output):
  # Allows for examination of inputs / outputs and modification of the outputs
  # after the forward pass. Note that inputs are always wrapped in a tuple while outputs
  # are passed as-is.

  # Residual computation a la ResNet.
  return output + inputs[0]

def backward_hook(m, grad_inputs, grad_outputs):
  # Allows for examination of grad_inputs / grad_outputs and modification of
  # grad_inputs used in the rest of the backwards pass. Note that grad_inputs and
  # grad_outputs are always wrapped in tuples.
  new_grad_inputs = [torch.ones_like(gi) * 42. for gi in grad_inputs]
  return new_grad_inputs

# Create sample module & input.
m = nn.Linear(3, 3)
x = torch.randn(2, 3, requires_grad=True)

# ==== Demonstrate forward hooks. ====
# Run input through module before and after adding hooks.
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>)

# Note that the modified input results in a different output.
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>)

# Note the modified output.
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>)

# Remove hooks; note that the output here matches the output before adding hooks.
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>)

# ==== Demonstrate backward hooks. ====
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]])

# Clear gradients before running backward pass again.
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., 42., 42.]])

高级功能

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

分布式训练

PyTorch 中存在各种分布式训练方法，这两种方法都用于使用多个 GPU 扩展训练 以及跨多台机器的训练。查看分布式训练概述页面 有关如何使用这些的详细信息。

性能分析

PyTorch Profiler 可用于识别 模型内的性能瓶颈。它测量并输出 内存使用情况和花费的时间。

通过 Quantization 提高性能

将量化技术应用于模块可以通过使用较低的 bitwidth 的 bitwidth 比浮点精度多。在此处查看 PyTorch 提供的各种量化机制。

通过修剪提高内存使用率

大型深度学习模型通常参数化过度，从而导致高内存使用率。为了解决这个问题，PyTorch 提供模型修剪机制，这有助于在保持任务准确性的同时减少内存使用量。修剪教程介绍了如何使用 PyTorch 提供的修剪技术，或根据需要定义自定义修剪技术。

参数化

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

使用 FX 转换模块

PyTorch 的 FX 组件提供了一种灵活的转换方式 modules 的 Module。这可用于以编程方式生成 或 操作模块，适用于各种用例。要探索 FX，请查看这些使用 FX 进行卷积 + 批量范数融合和 CPU 性能分析的示例。