pytorch模型结构可视化

最近在做跟hrnet相关的一些工作，需要查看网络的结构，对某些网络层需要进行特别的处理，所以需要对网络结构进行可视化。一开始想到的就是netron，但是由于hrnet网络中对层进行了封装处理，所以得到的是如下示意图：

但是这种对封装的模块不能展开，不方便查看内部结构。

后来尝试使用了TensorBoardX，代码大致如下：

# build model
model = eval('models.'+config.MODEL.NAME +
                 '.get_seg_model')(config)

dump_input = torch.rand(
        (1, 3, config.TRAIN.IMAGE_SIZE[1], config.TRAIN.IMAGE_SIZE[0]))
print("dump_input:", dump_input.shape)
# jing
with SummaryWriter(comment='hrnet') as w:
    w.add_graph(model, (dump_input,))

正常情况下，上面的代码运行结束后，会在当前目录生成一个叫runs的文件夹，runs文件夹里面会有一个文件夹Jul20_18-03-19_WH-PC19012HrNet，里面存储了可视化所需要的日志信息。用cmd进入到runs文件夹所在的目录中（路劲中不能有中文），然后cmd中输入：

tensorboard --logdir Jul_20-03-19_WH-PC19012HrNet

但是会报如下错误：

网上查了下可能跟网络结构中的None有关系，也可能跟pytorch TensorBoard版本有关系，反正没解决😅

后来使用graphviz+torchviz来可视化模型，主要步骤如下：

pip install graphviz  # 安装graphviz
pip install git+https://github.com/szagoruyko/pytorchviz  # 通过git安装torchviz

import torch
from torchvision.models import AlexNet
from torchviz import make_dot
 
x=torch.rand(8,3,256,512)
model=AlexNet()
y=model(x)

# 调用make_dot()函数构造图对象
g = make_dot(y)

# 保存模型，以PDF格式保存
g.render('Alex_model', view=False)