Netronは、onnxファイルの各レイヤーの形状メソッドを表示します

次のコードを使用してonnxモデルをエクスポートすると、Netronを使用して表示した結果は次のようになります。

import torch
import torch.onnx
import torchvision.models as models
from collections import OrderedDict

device = torch.device("cpu")

def convert():
    model = models.inception_v3(pretrained = False)
    pthfile = './inception_v3_google-1a9a5a14.pth'
    state_dict =torch.load(str(pthfile))
    new_state_dict = OrderedDict()
    
    for k, v in state_dict.items():
        if(k[0:7] == "module."):
          name = k[7:]
        else:
          name = k[0:]        
        new_state_dict[name] = v
    model.load_state_dict(new_state_dict)
    
    
    print(model)

    input_names = ["actual_input_1"]
    output_names = ["output1"]
    dummy_input = torch.randn(16, 3, 224, 224)
    dynamic_axes = {
    
    'actual_input_1': {
    
    0: '-1'}, 'output1': {
    
    0: '-1'}}

    torch.onnx.export(model, dummy_input, "inception_v3.onnx", input_names = input_names, output_names = output_names, dynamic_axes = dynamic_axes, opset_version=11)
    
if __name__ == "__main__":
    convert()

しかし、オペレーターによって出力された形状結果を表示したい場合があります。明らかに、上の図からそれを表示する方法はありません。次に、この時点で派生するために、onnxライブラリにshape_inferenceが必要です。コードは次のとおりです。

import onnx
from onnx import shape_inference
path = "..." #the path of your onnx model
onnx.save(onnx.shape_inference.infer_shapes(onnx.load(path)), path)

Netronを使用して、保存されたモデルの結果を表示します。

必要な効果を正常に達成します。！！