模块
PyTorch 使用模块来表示神经网络。模块包括:
-
状态化计算的构建块。 PyTorch 提供了丰富的模块库,并简化了新自定义模块的定义过程,从而使复杂多层神经网络的构建变得更加容易。
-
与 PyTorch 的自动微分 系统紧密集成。 模块使为 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
,以便与其它模块进行组合使用。 -
它定义了一些用于计算的状态。在这里,状态由随机初始化的
weight
和bias
张量组成,这些张量定义了仿射变换。因为每个都被定义为Parameter
,它们被注册到模块中,并且会自动从调用parameters()
时跟踪和返回。参数可以被认为是模块计算中的“可学习”部分(稍后会有更多介绍)。请注意,模块不一定需要有状态,也可以是无状态的。 -
它定义了一个执行计算的 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 的自动微分系统会自动处理这个后向传递计算,因此不需要为每个模块手动实现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
,然后再将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.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()
方法:
# 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>)
更一般地说,可以使用apply()
函数将任意函数递归地应用到一个模块及其所有子模块上。例如,要对模块及其子模块的参数进行自定义初始化:
# 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 提供了一个包含高性能模块的大库,在 torch.nn
命名空间中执行常见的神经网络操作,如池化、卷积和损失函数等。
在下一节中,我们将提供一个训练神经网络的完整示例。
如需了解更多信息,请参阅:
利用模块进行神经网络训练
一旦建立了网络,就需要对其进行训练,并且可以使用 PyTorch 中的优化器(如 torch.optim
提供的那些)轻松地优化其参数。
# 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() ...
在这个简化的例子中,网络学习简单地输出零,因为任何非零输出都会根据其绝对值被“惩罚”,这里使用torch.abs()
作为损失函数。尽管这不是一个非常有趣的任务,但训练的关键部分都已包含:
-
建立了一个网络。
-
创建了一个优化器(例如随机梯度下降优化器),并将其与网络的参数相关联。
-
- 训练循环
-
-
接收输入
-
运行网络
-
计算损失
-
将网络参数的梯度清零
-
调用 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: # 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])
训练神经网络常常会比较棘手。欲了解更多信息,请参阅:
模块状态
在上一节中,我们演示了如何训练模块的“参数”,即可学习的计算部分。现在,如果我们想将训练好的模型保存到磁盘,可以通过保存其 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>
模块的 state_dict
包含影响其计算的状态信息。这不仅包括模块的参数,还可能包含超出参数范围但会影响模块计算且不可学习的状态信息。为此,PyTorch 提供了“持久”和“非持久”缓冲区的概念。以下是模块可以拥有的各种状态类型的概述:
-
参数: 计算过程中可以学习的方面,存储在
state_dict
中 -
Buffers: 计算中的不可学习因素
-
持久化缓冲区:存储在
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]))])) # 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)
持久和非持久缓冲区都会受到使用to()
应用的整个模型范围内的设备/数据类型更改的影响:
# Moves all module parameters and buffers to the specified device / 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__() # 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
-
序列化语义:参考文档
-
什么是状态字典?参见详情
模块初始化
默认情况下,由torch.nn
提供的模块的参数和浮点缓冲区在模块实例化时,在CPU上以32位浮点值进行初始化,并采用历史上被认为适合该模块类型的方案。对于某些用例,可能希望使用不同的数据类型、设备(例如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)
请注意,上述展示的设备和数据类型选项也适用于模块中注册的所有浮点缓冲区。
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()
; 有关详细说明,请参阅此页面。
如需了解更多信息,请参阅:
模块钩子
在使用模块进行神经网络训练中,我们演示了模块的训练过程:迭代地执行前向和后向传递,并在每次迭代时更新模块参数。为了更精细地控制这一过程,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()
函数。
以下是一个展示正向和反向钩子用法的例子:
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 还提供了一些高级功能,旨在与模块配合使用。所有这些功能都适用于自定义编写的模块,但需要注意的是,某些功能可能需要模块符合特定约束条件才能支持。有关这些功能及其相应要求的详细讨论,请参见以下链接。
性能分析
PyTorch 分析器 可帮助您识别模型中的性能瓶颈。它会测量并报告内存使用和时间消耗等性能指标。
利用剪枝优化内存使用
大型深度学习模型通常过度参数化,导致内存使用量高。为了解决这个问题,PyTorch 提供了模型剪枝功能,这有助于在保持任务准确性的同时减少内存使用。PyTorch 剪枝教程 描述了如何利用 PyTorch 提供的剪枝技术或根据需要定义自定义剪枝方法。
参数化
对于某些应用程序,在模型训练期间限制参数空间是有益的。例如,强制学习到的参数保持正交性可以提高RNNs的收敛速度。PyTorch提供了应用参数化(如正交性)的机制,并进一步允许定义自定义约束条件。
模块的FX转换
PyTorch 的FX组件提供了一种灵活的方式来转换模块,通过直接操作模块的计算图。这可以用于程序化生成或修改模块,适用于各种应用场景。要了解 FX,请查看这些示例: 卷积+批量归一化融合 和 CPU性能分析。