首先安装graphviz
demo
可视化以下mlp
import torch
from torch import nn
from torchviz import make_dot
model = nn.Sequential()
model.add_module('W0', nn.Linear(8, 16))
model.add_module('tanh', nn.Tanh())
model.add_module('W1', nn.Linear(16, 1))
x = torch.randn(1, 8)
make_dot(model(x), params=dict(model.named_parameters())).view()
如下得到一个pdf文件