torchviz - Pytorch backward pass 시각화

pip install torchviz
apt-get install graphviz
# brew install graphviz -> for mac

from torchviz import make_dot

...
output = model(input)
loss = loss_fn(output)

graph = make_dot(loss, params=dict(model.named_parameters()))
graph.render(filename='file_name', format='png')  # graphviz install 필요, format default: pdf