Google Colaboratory 上で ONNX ファイルを Netron サーブしてブラウザ上でモデルを直接可視化する方法

import torch
import torch.nn as nn
import netron
import numpy as np
import onnx
from onnxsim import simplify
from google.colab import output

def view_netron(*, onnx_file_path: str):
  netron.start(onnx_file_path)
  output.serve_kernel_port_as_window(8081)

class ArgMaxModel(nn.Module):
    def __init__(
        self,
    ):
        super(ArgMaxModel, self).__init__()

    def forward(self, x):
        return torch.argmax(x, dim=1, keepdim=False)

OPSET = 11
MODEL_NAME = f'xxxx'
model = ArgMaxModel()
onnx_file = f'{MODEL_NAME}_{OPSET}.onnx'
x = torch.randn(1, 3, 128, 128)
torch.onnx.export(
    model,
    args=(x),
    f=onnx_file,
    opset_version=OPSET,
    input_names=[
        f'{MODEL_NAME}_input',
    ],
    output_names=[
        f'{MODEL_NAME}_output',
    ],
)
model_onnx1 = onnx.load(onnx_file)
model_onnx1 = onnx.shape_inference.infer_shapes(model_onnx1)
onnx.save(model_onnx1, onnx_file)
model_onnx2 = onnx.load(onnx_file)
model_simp, check = simplify(model_onnx2)
onnx.save(model_simp, onnx_file)

view_netron(onnx_file_path=onnx_file)

Serving 'xxxx.onnx' at http://localhost:18665
https://localhost:8081/