在深度学习论文、技术报告或教学分享中,一张清晰美观的神经网络结构图能让复杂的模型架构一目了然。但很多开发者苦于用Visio、PPT画的图不够专业,用专业绘图工具又门槛太高。本文将分享基于Python的神经网络可视化方案,从基础全连接层到复杂卷积网络,手把手教你画出期刊级别的神经网络图,附完整可运行代码!
一、核心工具选型
画神经网络图的核心需求是:代码易写、样式可定制、输出高清矢量图。对比主流工具后,推荐以下组合:
| 工具 | 优势 | 适用场景 |
|---|---|---|
| Matplotlib | Python生态标配,高度自定义,支持矢量图输出 | 全连接层、简单CNN/RNN可视化 |
| Graphviz + PyTorchviz | 自动解析模型计算图,无需手动布局 | 快速生成PyTorch模型结构 |
| NN-SVG | 在线工具(辅助) | 快速生成SVG格式的经典网络(LeNet/ResNet) |
本文重点讲解Matplotlib手动绘制(灵活度最高)和PyTorchviz自动生成(效率最高)两种方案。
二、基础版:Matplotlib绘制全连接神经网络
全连接层(MLP)是最基础的神经网络结构,先从它入手,掌握节点、连接、标注的绘制逻辑。
2.1 核心思路
- 定义各层节点数量(输入层、隐藏层、输出层);
- 计算每层节点的坐标位置(垂直居中,水平均分);
- 绘制节点(圆形)和连接(直线);
- 添加标注和样式美化。
2.2 完整代码
importmatplotlib.pyplotaspltimportnumpyasnp# 设置中文字体(解决中文乱码)plt.rcParams["font.family"]="SimHei"plt.rcParams["axes.unicode_minus"]=FalseclassNeuralNetworkPlot:def__init__(self,layers,figsize=(10,6),dpi=100):""" 初始化神经网络绘图类 :param layers: 列表,各层节点数,如[4, 8, 3]表示输入层4节点,隐藏层8节点,输出层3节点 :param figsize: 画布大小 :param dpi: 分辨率 """self.layers=layers self.n_layers=len(layers)self.fig,self.ax=plt.subplots(figsize=figsize,dpi=dpi)# 隐藏坐标轴self.ax.axis('off')def_get_node_coords(self):"""计算各层节点的坐标"""coords=[]fori,n_nodesinenumerate(self.layers):# 水平位置:均分x轴(0到1)x=i/(self.n_layers-1)ifself.n_layers>1else0.5# 垂直位置:均分y轴(0到1),居中对齐y=np.linspace(0,1,n_nodes)ifn_nodes%2==0:# 偶数节点微调居中y+=0.5/n_nodes coords.append(np.column_stack([np.full(n_nodes,x),y]))returncoordsdefdraw_connections(self,coords,line_color='#cccccc',line_width=1):"""绘制层间连接"""foriinrange(self.n_layers-1):fornode1incoords[i]:fornode2incoords[i+1]:self.ax.plot([node1[0],node2[0]],[node1[1],node2[1]],color=line_color,linewidth=line_width,zorder=1)defdraw_nodes(self,coords,node_size=200,node_color='#4285f4',edge_color='white'):"""绘制节点"""forlayer_coordsincoords:self.ax.scatter(layer_coords[:,0],layer_coords[:,1],s=node_size,c=node_color,edgecolors=edge_color,zorder=2,alpha=0.9)defadd_labels(self,coords,layer_labels=None,fontsize=12,fontcolor='#333333'):"""添加层标签"""iflayer_labelsisNone:layer_labels=[f'第{i+1}层'foriinrange(self.n_layers)]fori,(layer_coords,label)inenumerate(zip(coords,layer_labels)):# 标签位置在每层右侧x=layer_coords[0,0]+0.02y=0.5# 垂直居中self.ax.text(x,y,label,fontsize=fontsize,color=fontcolor,ha='left',va='center',weight='bold')defstyle_plot(self,title='全连接神经网络结构',title_fontsize=14,bg_color='white'):"""美化图表"""self.fig.patch.set_facecolor(bg_color)self.ax.set_title(title,fontsize=title_fontsize,pad=20,weight='bold')# 调整x/y轴范围,避免内容溢出self.ax.set_xlim(-0.05,1.05)self.ax.set_ylim(-0.05,1.05)defdraw(self,layer_labels=None):"""主绘制函数"""coords=self._get_node_coords()self.draw_connections(coords)self.draw_nodes(coords)self.add_labels(coords,layer_labels)self.style_plot()plt.tight_layout()plt.show()# 实例化并绘制(输入层4节点,隐藏层8节点,输出层3节点)nn_plot=NeuralNetworkPlot(layers=[4,8,3])nn_plot.draw(layer_labels=['输入层','隐藏层','输出层'])2.3 效果展示
运行代码后会生成如下效果:
- 输入层4个节点、隐藏层8个节点、输出层3个节点;
- 浅灰色连接线条,蓝色节点带白色边框,层次分明;
- 中文标签清晰,整体风格简洁专业;
- 输出为矢量图,可直接导出为PDF/PNG用于论文/报告。
2.4 定制化技巧
- 修改
node_color参数可更换节点颜色(如#ea4335红色、#fbbc05黄色); - 调整
node_size改变节点大小,适合不同场景; - 多层隐藏层只需修改
layers参数(如[5, 10, 10, 2]); - 导出高清图:将
plt.show()替换为plt.savefig('nn_plot.pdf', dpi=300, bbox_inches='tight')。
三、进阶版:PyTorchviz自动生成模型计算图
对于复杂的CNN/RNN/Transformer,手动绘制效率低,可使用PyTorchviz自动解析模型结构。
3.1 环境安装
pipinstalltorch torchviz graphviz注意:需额外安装Graphviz软件(https://graphviz.org/download/),并配置环境变量。
3.2 代码示例(以简单CNN为例)
importtorchimporttorch.nnasnnfromtorchvizimportmake_dot# 定义一个简单的CNN模型classSimpleCNN(nn.Module):def__init__(self):super(SimpleCNN,self).__init__()self.conv1=nn.Conv2d(3,16,kernel_size=3,stride=1,padding=1)self.relu=nn.ReLU()self.pool=nn.MaxPool2d(kernel_size=2,stride=2)self.fc1=nn.Linear(16*16*16,10)# 假设输入为3x32x32defforward(self,x):x=self.pool(self.relu(self.conv1(x)))x=x.view(x.size(0),-1)x=self.fc1(x)returnx# 实例化模型并生成输入model=SimpleCNN()x=torch.randn(1,3,32,32)# batch_size=1, channel=3, 32x32y=model(x)# 生成计算图dot=make_dot(y,params=dict(model.named_parameters()))# 保存为PDF(矢量图)或PNGdot.render("cnn_model",format="pdf")# 生成cnn_model.pdfdot.view()# 打开可视化窗口3.3 优势与适用场景
- 自动解析:无需手动定义节点/连接,直接基于模型代码生成;
- 支持复杂模型:CNN、RNN、Transformer、ResNet等均适用;
- 显示参数维度:可直观看到每层参数数量和张量维度;
- 适合快速生成模型架构图,用于技术文档或代码注释。
四、高级美化技巧
4.1 配色方案
- 专业期刊风格:主色
#2c3e50(深蓝灰)、辅助色#3498db(浅蓝); - 科技感风格:主色
#0077b6(科技蓝)、连接色#e0e0e0(浅灰); - 避免高饱和度颜色,保持视觉舒适。
4.2 细节优化
- 隐藏层节点可设置渐变颜色,区分不同层;
- 重要节点(如输出层)增大尺寸或更换颜色;
- 为连接线条添加透明度,避免密集线条重叠;
- 导出时选择矢量格式(PDF/SVG),保证放大后不失真。
4.3 混合方案
对于复杂模型,可结合两种方式:
- 用PyTorchviz生成基础结构;
- 导出为SVG后,用Inkscape/Figma手动美化(调整颜色、添加标注、简化冗余连接)。
五、常见问题解决
5.1 中文乱码
# 方案1:指定系统字体plt.rcParams["font.family"]=["SimHei","Arial Unicode MS","DejaVu Sans"]# 方案2:手动加载字体文件frommatplotlib.font_managerimportFontProperties font=FontProperties(fname='simhei.ttf',size=12)5.2 Graphviz报错
- 检查Graphviz是否安装并配置环境变量;
- Windows系统需重启终端/IDE,确保环境变量生效;
- 更换
dot.render的输出路径,避免中文路径。
5.3 节点重叠
- 调整
layers参数,合理分配每层节点数; - 增大画布尺寸(
figsize参数); - 调整节点垂直间距(修改
_get_node_coords中的y计算逻辑)。
六、总结
本文分享了两种实用的神经网络可视化方案:
- Matplotlib手动绘制:灵活度高,适合定制化需求,尤其适合全连接层、教学场景;
- PyTorchviz自动生成:效率高,适合复杂模型,快速生成计算图。
掌握这些方法后,你可以轻松画出符合论文、报告、教学需求的高颜值神经网络图。核心原则是:简洁清晰、层次分明、配色专业,避免过度装饰导致重点不突出。
如果需要绘制更复杂的网络(如Transformer的注意力层、ResNet的残差块),可以在Matplotlib基础上扩展节点形状(矩形、菱形)、添加箭头、标注维度信息等。希望本文能帮你摆脱“画图难”的困扰,让技术表达更专业。