[FIXED] Visualizing gradients in Pytorch models

Issue

I wanted to Visualize the gradients in each layer for each epoch in Pytorch Model. Any Library other than Tensorboard for this task?

Solution

Wandb.ai is a great tool for experiment tracking. It automatically logs the gradients and relevant system information if you call wandb.watch(model) on your model. It's also a good decentralized place to store the logs of all your runs and visualize them or compare them.

To give you the minimal code sample similar to the documentation which will log your loss, gradients and system info:

import wandb


wandb.init(project="your_project")
model = ... # set up your model

# Start tracking grads
wandb.watch(model, log='all')

for data, target in train_loader:
    output = model(data)
    loss = F.nll_loss(output, target)

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    wandb.log({"loss": loss})

Answered By - tnfru