How to plot models using torchviz or hiddenlayer

🀗Transformers
1

I am trying to plot models using torchviz and hiddenlayer but both gets errors.

torchviz - GitHub - waleedka/hiddenlayer: Neural network graphs and training metrics for PyTorch, Tensorflow, and Keras.
hiddenlayer - GitHub - szagoruyko/pytorchviz: A small package to create visualizations of PyTorch execution graphs

Common Code:

from transformers import AutoModel
model1 = AutoModel.from_pretrained("bert-base-uncased")

Code for torchviz :

from torchviz import make_dot, make_dot_from_trace
make_dot(model1)

Error:

AttributeError: ‘BertModel’ object has no attribute ‘size’

Code for hiddenlayer:

import hiddenlayer as hl
hl.build_graph(model1)

Error:

AssertionError: Argument args must be provided for Pytorch models.

How to solve these errors or do you have any other model plot method?

2

Hi,
There is a new package developed just for this purpose, called torchview. I tried your model with torchview with the following code

from transformers import AutoModel, AutoTokenizer
from torchview import draw_graph
model1 = AutoModel.from_pretrained("bert-base-uncased")

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
model = AutoModel.from_pretrained("bert-base-uncased")

inputs = tokenizer("Hello world!", return_tensors="pt")

model_graph = draw_graph(model1, input_data=inputs)

model_graph.visual_graph

The output is the following, (zoom in the image for clearity)

model
model1066×3773 360 KB

It also works very well with many other model types
The link for the project: torchview project
Disclaimer: I am the author of the package.

2 Likes
3

wow, thank you very much. I also try to build a library but I have no time so I skip it.

4

Hi Mert, thanks for your efforts in torchview, which is very useful to visualize models.

Though it works in some models (e.g. “bert-base-uncased” as you mentioned), there is a problem for some other models like “google/flan-t5-small”. Details can be found below

tokenizer = T5Tokenizer.from_pretrained("google/flan-t5-small")
model = T5ForConditionalGeneration.from_pretrained("google/flan-t5-small") 
input_text = "translate English to German: How old are you?"
inputs = tokenizer(input_text, return_tensors="pt")

model_graph = draw_graph(model, input_data=inputs).visual_graph
# Line below triggers an error: "RuntimeError: Failed to run torchgraph"
model_graph.render(directory='/home/jeff/tmp/graphviz_output').replace('\\', '/')