How to plot transformer model? (BERT Model)

[indramal]

I try to plot bert model using this library but I unable to do it.

from transformers import AutoModel, AutoTokenizer
model = AutoModel.from_pretrained("bert-base-uncased")
tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
inputs = tokenizer("Hello world!", return_tensors="pt")

from torchviz import make_dot, make_dot_from_trace
make_dot(model(inputs), params=dict(model.named_parameters()))

Error:

---

KeyError Traceback (most recent call last)
/usr/local/lib/python3.8/dist-packages/transformers/tokenization_utils_base.py in getattr(self, item)
247 try:
--> 248 return self.data[item]
249 except KeyError:

KeyError: 'size'

During handling of the above exception, another exception occurred:

AttributeError Traceback (most recent call last)
3 frames
/usr/local/lib/python3.8/dist-packages/transformers/tokenization_utils_base.py in getattr(self, item)
248 return self.data[item]
249 except KeyError:
--> 250 raise AttributeError
251
252 def getstate(self):

AttributeError:

[tanujjain]

Try replacing the following line:
make_dot(model(inputs), params=dict(model.named_parameters()))

with:

make_dot(model(inputs).mean(), params=dict(model.named_parameters()))
You could probably use some aggregation other than mean too.

[indramal]

I got this error for make_dot(model(inputs).mean(), params=dict(model.named_parameters()))

[Contributorrandom]

Use torchview package

from transformers import AutoModel, AutoTokenizer
from torchview import draw_graph

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

inputs = tokenizer("Hello world!", return_tensors="pt")
model_graph = draw_graph(model, input_data=inputs)
model_graph.visual_graph