ONNX ModelのGraph Outputsに既存のNode Outputを追加したい
6423 ワード
はじめに
PTQをONNX Runtimeとかで動かす際に,
各Conv NodeのOutputsを取り出したいなぁと思ったときの対応方法.
GraphSurgeonによるModel変換
ModelZooのresnet18を題材に,
resnetv15_batchnorm0_fwdというNodeをGraph Outputsに設定する.
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
Author And Source
この問題について(ONNX ModelのGraph Outputsに既存のNode Outputを追加したい), 我々は、より多くの情報をここで見つけました https://zenn.dev/rain_squallman/articles/fd56d08f30c916著者帰属:元の著者の情報は、元のURLに含まれています。著作権は原作者に属する。
Collection and Share based on the CC protocol