(1)首先安裝graphviz
(2)命令行安裝pytorchviz
pip install git+https://github.com/szagoruyko/pytorchviz
(3)示例代碼
import torch
import torch.nn as nn
import torchviz
from torch.autograd import Variable
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 = Variable(torch.randn(1,8))
y = model(x)
torchviz.make_dot(y.mean(), params=dict(model.named_parameters()))
參考文獻: