将 PyTorch Model 用可视化方法浏览 torchview，onxx, netron, summary | 撰写论文 paper

使用 PyTorch 构建的 Model，想要查看网络的图形，那么有以下方法，最终可视化效果最好的是：torchview； 对于优秀论文中的神经网络图则主要是通过 PPT 等更自由的工具手绘的¹。

torchview

文档

https://mert-kurttutan.github.io/torchview/tutorial/example_introduction/

需要安装：

• graphviz, 并保证 dot.exe 在 PATH 路径
• torchview

代码

def export_model_graph(model, input_sample, directory, filename = "model_graph", format = "svg", scale=5.0):
    '''
    Export model as image with torchview
    https://mert-kurttutan.github.io/torchview/reference/torchview/#torchview.torchview.draw_graph
    '''
    model_graph = draw_graph(model, 
                             input_size=input_sample.shape, 
                             expand_nested=True, 
                             directory=directory)
    model_graph.visual_graph.node_attr["fontname"] = "Helvetica"
    model_graph.resize_graph(scale=scale) # scale as per the view
    # https://graphviz.readthedocs.io/en/stable/api.html#graphviz.Graph
    model_graph.visual_graph.render(filename=filename, directory=directory, format=format)
    filepath = os.path.join(directory, "%s.%s" % (filename, format))

    if not os.path.exists(filepath):
        raise BaseException("Error on export torchview graph %s" % filepath)

    return filepath

图像示例

特点

清晰直观，信息丰富。

onnx, netron

https://onnx.ai/
https://github.com/lutzroeder/netron/tree/

导出 model 的 onnx 文件

def export_onnx_archive(model, filepath, input_sample):
    '''
    Export model as onnx format file
    https://pytorch.org/docs/stable/onnx.html
    '''
    is_training = False

    if model.training:
        # view the network graph in onnx format
        # https://pytorch.org/docs/stable/onnx.html
        model.eval()
        is_training = True

    torch.onnx.export(
        model,        # model to export
        (input_sample,),      # inputs of the model,
        filepath,        # filename of the ONNX model
        input_names=["input"],  # Rename inputs for the ONNX model
        dynamo=True             # True or False to select the exporter to use
    )
    
    if not os.path.exists(filepath):
        raise BaseException("File %s not found" % filepath)

    if is_training:
        model.train()

示例

特点

识别更多模型的信息。

缺点：

• 需要额外导出 onnx model
• 并没有 torchview 直观

torchinfo summary

from torchinfo import summary
logger.info(summary(resnet_model_50, verbose=0))

图片示例

特点

是文本的形式，简洁依赖少。

缺点：

• 表现力不如 torchview

nn-svg

生成 FCNN，AlexNet, LeNet 论文中，手绘风格的神经网络图片。可以手动设置参数。

使用

https://www.chatopera.com/files/nn-graph/index.html

开源项目

https://github.com/alexlenail/NN-SVG

特点

• 生成和论文中类似的风格的图片

---

1. https://datascience.stackexchange.com/questions/66343/drawing-neural-network-diagram-for-academic-papers ↩︎