Zenn
🐡

ONNX ModelのGraph Outputsに既存のNode Outputを追加したい

2022/04/24に公開
ディープラーニング
ONNX
tech

はじめに

PTQをONNX Runtimeとかで動かす際に,
各Conv NodeのOutputsを取り出したいなぁと思ったときの対応方法.

GraphSurgeonによるModel変換

ModelZooのresnet18を題材に,
resnetv15_batchnorm0_fwdというNodeをGraph Outputsに設定する.

GitHub Gist

import onnx
from onnx import shape_inference
import onnx_graphsurgeon as gs
from argparse import ArgumentParser
from typing import List

def add_output_node(orgmodel: str, modmodel: str, outlist: List[str]) -> int:
    model = shape_inference.infer_shapes(onnx.load(orgmodel))
    graph = gs.import_onnx(model)
    for outnode in outlist:
        tensor = graph.tensors()[outnode]
        graph.outputs.append(tensor)

    onnx.save(gs.export_onnx(graph), f'{modmodel}')

if __name__=='__main__':
    parser = ArgumentParser(description='add your specified node to graph output')
    parser.add_argument('orgmodel', help='import original model name')
    parser.add_argument('modmodel', help='export original model name')
    parser.add_argument('--nodes', nargs='*', required=True)
    args = parser.parse_args()

    ret = add_output_node(args.orgmodel, args.modmodel, args.nodes)

$ python3 append_output_node.py resnet18-v1-7.onnx resnet18-v1-7_test.onnx --nodes resnetv15_batchnorm0_fwd

append outputs

Discussion