Pytorch中用Tensorboard可视化网络结构

第一步——Import

from tensorboardX import SummaryWriter

第二步——“打一个log”

writer = SummaryWriter('log') 
#建立一个保存数据用的东西

第三步——网络实例化和测试

model = resnet(depth=19, dataset='cifar10')
#resnet是我自己定义的网络
dummy_input = torch.rand(1, 3, 32, 32).cuda()  
#（batch_size, channel, kernel, kernel）
# 记得张量后面要紧跟 .cuda() !!!!!!!
with SummaryWriter(comment='resnet') as w:
w.add_graph(model, (dummy_input,))

     强调一点： 在一般的代码中，神经网络都是在GPU上跑的，所以一定要加“ .cuda() ”。测试图片的尺寸的话，就直接搜索数据集简介好了。
     运行之后就会在runs文件夹下生成文件。

第四步——可视化

     在命令行中cd到runs文件夹的上层目录：

    输入命令：

tensorboard --logdir ./runs 

    接着直接用浏览器打开这个网址就好了。（出不来就刷新一下）