评价此页
★ ★ ★ ★ ★
intermediate/visualizing_gradients_tutorial

注意

跳转至页面底部 下载完整示例代码。

可视化梯度#

作者: Justin Silver

本教程解释了如何提取并可视化神经网络中任意层的梯度。通过检查信息如何从网络末端流向我们需要优化的参数，我们可以调试训练过程中出现的诸如 梯度消失或梯度爆炸 等问题。

开始之前，请确保您了解 张量及其操作方法。对 自动求导（autograd）原理 的基本了解也将有所帮助。

设置#

首先，请确保已 安装 PyTorch，然后导入必要的库。

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import matplotlib.pyplot as plt

接下来，我们将创建一个用于 MNIST 数据集的网络，类似于 批量归一化（Batch Normalization）论文 中描述的架构。

为了说明梯度可视化的重要性，我们将实例化两个版本的网络：一个带有批量归一化（BatchNorm），另一个不带。批量归一化是解决 梯度消失/爆炸 的一种极其有效的技术，我们将通过实验来验证这一点。

我们使用的模型具有可配置数量的重复全连接层，这些层在 nn.Linear、norm_layer 和 nn.Sigmoid 之间交替。如果启用了批量归一化，norm_layer 将使用 BatchNorm1d，否则它将使用 Identity（恒等映射）转换。

def fc_layer(in_size, out_size, norm_layer):
    """Return a stack of linear->norm->sigmoid layers"""
    return nn.Sequential(nn.Linear(in_size, out_size), norm_layer(out_size), nn.Sigmoid())

class Net(nn.Module):
    """Define a network that has num_layers of linear->norm->sigmoid transformations"""
    def __init__(self, in_size=28*28, hidden_size=128,
                 out_size=10, num_layers=3, batchnorm=False):
        super().__init__()
        if batchnorm is False:
            norm_layer = nn.Identity
        else:
            norm_layer = nn.BatchNorm1d

        layers = []
        layers.append(fc_layer(in_size, hidden_size, norm_layer))

        for i in range(num_layers-1):
            layers.append(fc_layer(hidden_size, hidden_size, norm_layer))

        layers.append(nn.Linear(hidden_size, out_size))

        self.layers = nn.Sequential(*layers)

    def forward(self, x):
        x = torch.flatten(x, 1)
        return self.layers(x)

接下来我们设置一些虚拟数据，实例化两个版本的模型，并初始化优化器。

# set up dummy data
x = torch.randn(10, 28, 28)
y = torch.randint(10, (10, ))

# init model
model_bn = Net(batchnorm=True, num_layers=3)
model_nobn = Net(batchnorm=False, num_layers=3)

model_bn.train()
model_nobn.train()

optimizer_bn = optim.SGD(model_bn.parameters(), lr=0.01, momentum=0.9)
optimizer_nobn = optim.SGD(model_nobn.parameters(), lr=0.01, momentum=0.9)

我们可以通过探测其中一个内部层来验证批量归一化是否仅应用于其中一个模型。

print(model_bn.layers[0])
print(model_nobn.layers[0])

Sequential(
  (0): Linear(in_features=784, out_features=128, bias=True)
  (1): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, bias=True, track_running_stats=True)
  (2): Sigmoid()
)
Sequential(
  (0): Linear(in_features=784, out_features=128, bias=True)
  (1): Identity()
  (2): Sigmoid()
)

注册钩子（Hooks）#

由于我们将模型的逻辑和状态封装在 nn.Module 中，如果我们想避免直接修改模块代码，就需要另一种方法来访问中间梯度。这可以通过 注册一个钩子（hook） 来实现。

警告

比起在张量本身上使用 retain_grad()，更推荐使用附加到输出张量上的反向传播钩子。另一种方法是直接附加模块钩子（例如 register_full_backward_hook()），前提是 nn.Module 实例没有执行任何原地（in-place）操作。有关详细信息，请参阅 此问题跟踪。

以下代码定义了我们的钩子并收集了网络各层的描述性名称。

# note that wrapper functions are used for Python closure
# so that we can pass arguments.

def hook_forward(module_name, grads, hook_backward):
    def hook(module, args, output):
        """Forward pass hook which attaches backward pass hooks to intermediate tensors"""
        output.register_hook(hook_backward(module_name, grads))
    return hook

def hook_backward(module_name, grads):
    def hook(grad):
        """Backward pass hook which appends gradients"""
        grads.append((module_name, grad))
    return hook

def get_all_layers(model, hook_forward, hook_backward):
    """Register forward pass hook (which registers a backward hook) to model outputs

    Returns:
        - layers: a dict with keys as layer/module and values as layer/module names
                  e.g. layers[nn.Conv2d] = layer1.0.conv1
        - grads: a list of tuples with module name and tensor output gradient
                 e.g. grads[0] == (layer1.0.conv1, tensor.Torch(...))
    """
    layers = dict()
    grads = []
    for name, layer in model.named_modules():
        # skip Sequential and/or wrapper modules
        if any(layer.children()) is False:
            layers[layer] = name
            layer.register_forward_hook(hook_forward(name, grads, hook_backward))
    return layers, grads

# register hooks
layers_bn, grads_bn = get_all_layers(model_bn, hook_forward, hook_backward)
layers_nobn, grads_nobn = get_all_layers(model_nobn, hook_forward, hook_backward)

训练与可视化#

现在让我们训练模型几个 epoch。

epochs = 10

for epoch in range(epochs):

    # important to clear, because we append to
    # outputs everytime we do a forward pass
    grads_bn.clear()
    grads_nobn.clear()

    optimizer_bn.zero_grad()
    optimizer_nobn.zero_grad()

    y_pred_bn = model_bn(x)
    y_pred_nobn = model_nobn(x)

    loss_bn = F.cross_entropy(y_pred_bn, y)
    loss_nobn = F.cross_entropy(y_pred_nobn, y)

    loss_bn.backward()
    loss_nobn.backward()

    optimizer_bn.step()
    optimizer_nobn.step()

在运行前向和反向传播后，所有中间张量的梯度应该会出现在 grads_bn 和 grads_nobn 中。我们计算每个梯度矩阵的平均绝对值，以便比较这两个模型。

def get_grads(grads):
    layer_idx = []
    avg_grads = []
    for idx, (name, grad) in enumerate(grads):
        if grad is not None:
            avg_grad = grad.abs().mean()
            avg_grads.append(avg_grad)
            # idx is backwards since we appended in backward pass
            layer_idx.append(len(grads) - 1 - idx)
    return layer_idx, avg_grads

layer_idx_bn, avg_grads_bn = get_grads(grads_bn)
layer_idx_nobn, avg_grads_nobn = get_grads(grads_nobn)

计算出平均梯度后，我们现在可以绘制它们，并观察数值如何随网络深度变化。请注意，当我们不使用批量归一化时，中间层的梯度值会非常迅速地降为零。然而，带有批量归一化的模型在其中间层保持了非零的梯度。

fig, ax = plt.subplots()
ax.plot(layer_idx_bn, avg_grads_bn, label="With BatchNorm", marker="o")
ax.plot(layer_idx_nobn, avg_grads_nobn, label="Without BatchNorm", marker="x")
ax.set_xlabel("Layer depth")
ax.set_ylabel("Average gradient")
ax.set_title("Gradient flow")
ax.grid(True)
ax.legend()
plt.show()
Gradient flow

结论#

在本教程中，我们演示了如何可视化封装在 nn.Module 类中的神经网络的梯度流。我们定性地展示了批量归一化如何帮助缓解深度神经网络中出现的梯度消失问题。

如果您想了解更多关于 PyTorch 自动求导系统的工作原理，请访问下方的 参考资料。如果您对本教程有任何反馈（改进建议、纠错等），请使用 PyTorch 论坛 和/或 问题跟踪器 与我们联系。

(可选) 附加练习#

  • 尝试增加模型中的层数（num_layers），看看这对梯度流图有什么影响。

  • 您将如何修改代码以可视化平均激活值（Activations）而不是平均梯度？（提示：在 hook_forward() 函数中，我们可以访问原始张量输出）

  • 还有哪些处理梯度消失和梯度爆炸的方法？

参考文献#

脚本总运行时间: (0 分钟 0.281 秒)

在此页面上
PyTorch 库