GPTのモデル構造を可視化した

GPTのモデル構造を目で見てみたい!
そんな気持ち、わかるでしょ?
技研フリマをやりながら、どうにかこうにか出力したよ

ご覧あれ

やり方メモ

from transformers import AutoTokenizer, AutoModelForCausalLM
from torchviz import make_dot
tokenizer = AutoTokenizer.from_pretrained("gpt2")
from transformers import pipeline, set_seed
generator = pipeline('text-generation', model='gpt2')
m= generator.model
x= m.generate()
y= m.forward(x)
image = make_dot(y.logits, params=dict(generator.model.named_parameters()), show_attrs=True)
image.format = "png"
image.render("NeuralNet")

単純に見えますがリバースエンジニアリング(リフレクション)しながら試行錯誤したのでこのコード書くのに1時間くらいかかってしまった・・・
ちなみに、show_savedオプションをつけるとPNGで出力できなくなってしまう(PDFは可能)。

いやこれは楽しいね。
GPT-2だけど(GPT3は非公開だから)
まあでも構造に大差ないはずなのでこれを印刷してポスターにでもするか