【Tools】神经网络、深度学习和机器学习模型可视化工具——Netron

背景

有时候我们写完深度学习模型后，想看看代码实现的模型和我们预期是否一致，但是没有一个好的工具。最近发现有一个软件Netron支持对一部分深度学习模型可视化，源码地址。Netron 支持 ONNX、TensorFlow Lite、Caffe、Keras、Darknet、PaddlePaddle、ncnn、MNN、Core ML、RKNN、MXNet、MindSpore Lite、TNN、Barracuda、Tengine、CNTK、TensorFlow.js、Caffe2 和 UFF。

Netron 为 PyTorch、TensorFlow、TorchScript、OpenVINO、Torch、Vitis AI、kmodel、Arm NN、BigDL、Chainer、Deeplearning4j、MediaPipe、ML.NET 和 scikit-learn 提供实验性支持。

手上有现成的基于pytorch实现的textcnn（模型很简单）代码，下面演示如下。

模型代码实现

需要说明的是，想可视化Pytorch实现的模型，需要将模型导出为torchscript文件。如果不了解torchrscript，可以参考pytorch官方文档->torchscript。代码实现如下：

import torch.nn as nn
import torch


class TextCNN(nn.Module):

    def __init__(self):
        super().__init__()
        vocab_size = 10000
        embedding_dim = 128
        filter_number = 60
        output_size = 10
        kernel_list = [3, 4, 5]
        max_length = 200
        drop_out = 0.5
        self.embedding = nn.Embedding(vocab_size, embedding_dim)
        self.convs = nn.ModuleList([
            nn.Sequential(nn.Conv1d(in_channels=embedding_dim, out_channels=filter_number,
                                    kernel_size=kernel),
                          nn.LeakyReLU(),
                          nn.MaxPool1d(kernel_size=max_length - kernel + 1))
            for kernel in kernel_list
        ])
        self.fc = nn.Linear(filter_number * len(kernel_list), output_size)
        self.dropout = nn.Dropout(drop_out)

    def forward(self, x) :
        x = self.embedding(x)
        x = x.permute(0, 2, 1)  # 维度转换，满足一维卷积的输入
        out = [conv(x) for conv in self.convs]
        out = torch.cat(out, dim=1)  # [128, 300, 1, 1]，各通道的数据拼接在一起
        out = out.view(x.size(0), -1)  # 展平
        out = self.fc(out)  # 结果输出[128, 2]
        out = self.dropout(out)  # 构建dropout层out = self.dropout(out)  # 构建dropout层
        return out

# 模型出初始化
model = TextCNN()

jit_script_model = torch.jit.script(model)
# torchscript模型导出
jit_script_model.save("model.pt")

模型可视化

根据自己的环境下载软件，并安装后，直接将生成的模型拖到界面中即可。如下图：

那么这时我们就可以就我们的模型可视化了。当然对应的模块，还会显示该模块的输入与输出以及其他信息，如下图所示：

之前也尝试查看bilstm+crf实现的模型，也是能够查看的。其他的模型我们可以根据自己的需要去试一试哦。