ONX 모델의Graph Outputs에 기존 노드 Output 추가

개시하다

ONX Runtime 등을 사용하여 PTQ를 시작할 때
각 Conv Node의 Outputs를 제거하려는 경우 해당되는 방법입니다.

GraphiSurgen을 통한 모델 변환

ModelZoo의 resnet 18를 소재로
resnetv15_batchnorm0_fwd Node를 Graph Outputs로 설정합니다.
!
shape_inference를 통해 tensor의 정보를 얻어야 하기 때문에
중간에 Inference가 없는 노드가 끼어 있을 때는 사용할 수 없습니다.
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