Is there any function in pytorch like plot_model function in the tensor flow to print the network architecture?

vision
1

I want to print the network architecture like the one shown in the figure. Is there any function in pytorch which can do so?

Capture
Capture561×541 7.66 KB

2

You can use torchviz library for the same!

3

This is my code which throwing errors. Can you provide me some more reference about the torchviz?

!pip install git+https://github.com/szagoruyko/pytorchviz
from torchviz import make_dot
make_dot(mlp,mlp.parameters())

4

Take a look at this issue:https://github.com/szagoruyko/pytorchviz/issues/74

5

Another library is torchview, which is very close to plot_model of keras as it can capture module hierarchy. The similarity to plot_model API was a big factor in the design of the library
For instance, output for mlp model is the following
mlp

Disclaimer: I am the author of library