Visualização do modelo de rede neural PyTorch (Netron)

Visualização do modelo de rede neural PyTorch (Netron)

O Netron é uma ferramenta para visualizar modelos de aprendizado profundo, o que pode nos ajudar a entender melhor a estrutura e os parâmetros do modelo.

Arquivos de armazenamento de modelo nos seguintes formatos são suportados:

Formatar	modelo (arquivo)	Abrir sem baixar
ONNX	squeezenet	abrir
TensorFlow Lite	yamnet	abrir
TensorFlow	robô de xadrez	abrir
Duro	rede móvel	abrir
TorchScriptName	traced_online_pred_layer	abrir
Core ML	exercitar	abrir
Darknet	yolo	abrir

Link do GitHub: https://github.com/lutzroeder/netron

Site oficial: https://netron.app

---

ONNX

(1) No PyTorch, você pode usar torch.onnx.exporta função para exportar o modelo para o formato ONNX:

import torch
import netron


# 定义 PyTorch 模型
class MyModel(torch.nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.conv = torch.nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.bn = torch.nn.BatchNorm2d(64)
        self.relu = torch.nn.ReLU(inplace=True)
        self.pool = torch.nn.MaxPool2d(kernel_size=2, stride=2)
        self.fc = torch.nn.Linear(64 * 8 * 8, 10)

    def forward(self, x):
        x = self.conv(x)
        x = self.bn(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(-1, 64 * 8 * 8)
        x = self.fc(x)
        return x


# 创建模型实例并加载预训练权重
model = MyModel()

# 设置示例输入
input = torch.randn(1, 3, 32, 32)

# 将模型导出为 ONNX 格式
torch.onnx.export(model, input, './model/Test/onnx_model.onnx')  # 导出后 netron.start(path) 打开

netron.start(2) Abra o arquivo de modelo ONNX exportado com o comando Netron :

import netron

# 打开导出的 ONNX 模型文件
netron.start('./model/Test/onnx_model.onnx')

Serving './model/Test/onnx_model.onnx' at http://localhost:8080

A ferramenta Netron será iniciada automaticamente no navegador e visualizará o arquivo de modelo.

Perceber:

Quando o modelo for exportado para o formato ONNX, um arquivo .onnxcom , e também pode ser visualizado ao carregá-lo no site oficial da Netron :

No Netron, você pode visualizar a estrutura do modelo, parâmetros, entrada e saída e outras informações. A visualização do modelo pode ser ajustada ampliando, girando e traduzindo para entender melhor a estrutura e os parâmetros do modelo.

tocha.salvar

Ao torch.savevisualizar um modelo salvo com:

# 保存模型
torch.save(model.state_dict(), './model/Test/saved_model.pt')

# 可视化
netron.start('./model/Test/saved_model.pt')

Conforme mostrado na figura abaixo, este método não exibe as informações detalhadas do modelo:

Portanto: o Netron não oferece suporte a torch.savearquivos de modelo exportados pelo PyTorch.

tocha.jit.script

Para referência: maçarico.jit.script e maçarico.jit.trace

Use para torch.jit.scriptconverter o modelo em um script primeiro e, em seguida, torch.jit.savesalve o modelo para visualizá-lo:

# TorchScript：script
scripted_model = torch.jit.script(model)

# 保存模型
torch.jit.save(scripted_model, './model/Test/scripted_model.pth')

# 可视化
netron.start('./model/Test/scripted_model.pth')

tocha.jit.trace

Para referência: maçarico.jit.script e maçarico.jit.trace

Use para torch.jit.traceprimeiro converter o modelo em uma ferramenta que rastreia a execução do modelo, depois use torch.jit.savepara salvar o modelo e, finalmente, visualizá-lo:

# TorchScript：trace
traced_model = torch.jit.trace(model, torch.randn(1, 3, 32, 32))

# 保存模型
torch.jit.save(traced_model, './model/Test/traced_model.pth')

# 可视化
netron.start('./model/Test/traced_model.pth')