通过 Tensorboard 检测 Pytorch 中梯度消失/爆炸的最佳方法

Question

我怀疑我的 Pytorch 模型梯度消失。我知道我可以跟踪每一层的梯度并用writer.add_scalar或记录它们writer.add_histogram。然而，对于具有相对较多层数的模型，在 TensorBoard 日志上显示所有这些直方图和图表会变得有点麻烦。我并不是说它不起作用，只是每个图层都有不同的图表和直方图并滚动它们有点不方便。

我正在寻找一个图表，其中y轴（垂直）表示梯度值（特定层的梯度平均值），轴x（水平）显示层数（例如，at 的值x=1是第一层的梯度值），轴z（深度）是纪元数。

这看起来像直方图，但当然，它与直方图有本质上的不同，因为轴x不代表 beans。人们可以编写一段肮脏的代码来创建一个直方图，其中代替 bean 的是层编号，类似于（显然，这是伪代码）：

fake_distribution = []
for i, layer in enumerate(model.layers):
   fake_distribution += [i for j in range(int(layer.grad.mean()))]
writer.add_histogram('gradients', fake_distribution)

我想知道是否有更好的方法。

Answer 1

这是一个最小的示例，说明如何评估模型中特定层的范数。采用一个简单的模型进行说明：

class ConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 10, 5)
        self.conv2 = nn.Conv2d(10, 20, 5)
        self.fc1 = nn.Linear(8000, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, input):
        x = F.relu(self.conv1(input))
        x = F.relu(self.conv2(x))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return x

net = ConvNet()
net(torch.rand(5,1,28,28)).mean().backward()

看看clip_grad_norm_作为参考。要测量层上梯度的大小，conv1您可以：计算由属于该层的参数的 L2 梯度范数组成的向量的 L2 范数。这是通过以下代码完成的：

parameters = net.conv1.parameters()
norm_type = 2
total_norm = torch.norm(
    torch.stack([torch.norm(p.grad.detach(), norm_type) for p in parameters]), norm_type)

或者，您可以取该层上最大梯度分量的最大值，即inf-norm：

total_norm = torch.max(
     torch.stack([p.grad.detach().abs().max() for p in parameters]))

要将它们登录到您的 TensorBoard 上，您可以add_scalar在您的SummaryWriter：

for name, module in net.named_children():
    norm = torch.norm(
        torch.stack([torch.norm(p.grad.detach(), 2) for p in parameters]), 2)
    writer.add_scalar(f'check_info/{name}', norm, iter)