pytorch模型可视化,torchviz,tensorboardX,文本方式

torchviz生成一个pdf,pdf怎样命名还不知道,或许只能默认命名。

注意,torchviz我试了下将模型和伪造输入移到cuda上就会报错,所以这种可视化直接在cpu运行上即可。

torchviz看梯度类型更清楚,tensorboard看数据流形状更清楚。

还有tensorboardX方式,两种方式代码:

 1 from tensorboardX import SummaryWriter
 2 from torchviz import make_dot
 3 
 4 # 准备数据
 5 inputs_fake = torch.randn(13, 1, 32, 32)
 6 model = glimpse_cnn  # 模型实例
 7 
 8 # TensorboardX方式
 9 with SummaryWriter(comment='vgg') as w:
10     w.add_graph(model, (inputs_fake,))
11 
12 # torchviz方式. inputs_fake后面跟.requires_grad_(True)显示x输入形状框
13 vis_graph = make_dot(model(inputs_fake.requires_grad_(True)),
14                      params=dict(list(model.named_parameters()) + [('x', inputs_fake)]))
15 vis_graph.view()

也可以直接保存为文本格式,很详细,方便后面查看:

1 with open('model.txt', 'w') as w:
2     w.write(str(model))