安装graphviz
pip install graphviz
sudo apt install graphviz
dot -version
若安装成功则会显示:
Using device: dot:dot:core
The plugin configuration file:
/usr/lib/x86_64-linux-gnu/graphviz/config6a
was successfully loaded.
render : cairo dot dot_json fig gd json json0 map mp pic pov ps svg tk vml vrml xdot xdot_json
layout : circo dot fdp neato nop nop1 nop2 osage patchwork sfdp twopi
textlayout : textlayout
device : canon cmap cmapx cmapx_np dot dot_json eps fig gd gd2 gif gv imap imap_np ismap jpe jpeg jpg json json0 mp pdf pic plain plain-ext png pov ps ps2 svg svgz tk vml vmlz vrml wbmp x11 xdot xdot1.2 xdot1.4 xdot_json xlib
loadimage : (lib) eps gd gd2 gif jpe jpeg jpg png ps svg xbm
添加环境变量
export PATH=$PATH:/usr/lib/x86_64-linux-gnu/graphviz
生成网络图片
运行下面代码
import torch
import torch.nn as nn
from torch.autograd import Variable
from torchviz import make_dot
class LeNet(nn.Module):
def __init__(self):
super(LeNet