pytorch模型可視化，torchviz，tensorboardX，文本方式

torchviz生成一個pdf，pdf怎樣命名還不知道，或許只能默認命名。

注意，torchviz我試了下將模型和偽造輸入移到cuda上就會報錯，所以這種可視化直接在cpu運行上即可。

還有tensorboardX方式，兩種方式代碼：

from tensorboardX import SummaryWriter
from torchviz import make_dot

# 准備數據
inputs_fake = torch.randn(13, 1, 32, 32)
model = glimpse_cnn  # 模型實例

# TensorboardX方式
with SummaryWriter(comment='vgg') as w:
    w.add_graph(model, (inputs_fake,))

# torchviz方式. inputs_fake后面跟.requires_grad_(True)顯示x輸入形狀框
vis_graph = make_dot(model(inputs_fake.requires_grad_(True)),
                     params=dict(list(model.named_parameters()) + [('x', inputs_fake)]))
vis_graph.view()

 

也可以直接保存為文本格式，很詳細，方便后面查看:

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