Best way to detect Vanishing/Exploding gradient in Pytorch via Tensorboard

user_6411761 0
Viewed 2075

I suspect my Pytorch model has vanishing gradients. I know I can track the gradients of each layer and record them with writer.add_scalar or writer.add_histogram. However, with a model with a relatively large number of layers, having all these histograms and graphs on the TensorBoard log becomes a bit of a nuisance. I'm not saying it doesn't work, it's just a bit inconvenient to have different graphs and histograms for each layer and scroll through them.

I'm looking for a graph where the y axis (vertical) represents the gradient value (mean of gradient of a specific layer), the x axis (horizontal) shows the layer number (e.g. the value at x=1 is the gradient value for 1st layer), and the z axis (depth) is the epoch number.

This would look like a histogram, but of course, it would be essentially different from a histogram since the x axis does not represent beans. One can write a dirty code that would create a histogram where instead of beans there would be layer numbers, something like (this is a pseudo-code, obviously):

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)

I was wondering if there is a better way for this.

pytorch tensorboard deep-learning
1 Answers

This is a minimal example of how you could go about evaluating the norm of a particular layer in your model. Taking a simple model for illustration purposes:

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

Looking at clip_grad_norm_ as reference. To measure the magnitude of the gradient on layer conv1 you could: compute the L2-norm of the vector comprised of the L2-gradient-norms of parameters belonging to that layer. This is done with the following code:

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)

Alternatively, you can take the maximum of maximum gradient component on that layer i.e. the inf-norm:

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

To log them onto your TensorBoard, you can use add_scalar on your 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)
user_63313690
Related