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

Question

我怀疑我的 Pytorch 模型的梯度消失了。我知道我可以跟踪每一层的渐变并用writer.add_scalar 或writer.add_histogram 记录它们。但是，对于具有相对大量层数的模型，在 TensorBoard 日志中包含所有这些直方图和图形变得有点麻烦。我并不是说它不起作用，只是为每一层设置不同的图形和直方图并滚动浏览它们有点不方便。

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

这看起来像一个直方图，但当然，它与直方图有本质的不同，因为x 轴不代表豆类。可以编写一个脏代码，它会创建一个直方图，而不是 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，您可以在您的 SummaryWriter 上使用 add_scalar：

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)