Pytorch下网络结构可视化

requirements

torch
tensorboard
tensorboardX

其中，tensorboardX必须与tensorboard一起安装到同一环境下。

使用方法

import torch
from tensorboardX import SummaryWriter

model = LeNet5(4,2) # 你的模型在此处生成实例
x = torch.rand(1,4,28,28) # 随机生成一个 batch_size为 1 的输入变量
with SummaryWriter(comment='LeNet5') as w:
    w.add_graph(model,x)

运行上面的代码，将会在文件所在目录生成一个 runs 文件夹，其下有一个名字类似于：Nov03_17-07-26_DESKTOP-1CSBK3QLeNet5的文件夹，复制该路径（相对路径或者绝对路径都可以，路径中不可以有中文），然后在 cmd 中输入：

tensorboard --logdir runs\Nov03_17-07-26_DESKTOP-1CSBK3QLeNet5

回车后，你将看到：

TensorFlow installation not found - running with reduced feature set.
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.3.0 at http://localhost:6006/ (Press CTRL+C to quit)

然后打开这个链接即可。

Pytorch下网络结构可视化

需要注意的

runs的目录，一定要是相对路径或者绝对路径，不要搞错，否则读取不到文件，打开链接就是一个报错的界面。

参考资料

[1] pytorch 网络结构可视化方法汇总（三种实现方法详解