pytorch可视化网络结构

首先安装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文件