美文网首页
onnx warn消除

onnx warn消除

作者: 小天奇 | 来源:发表于2022-11-25 14:40 被阅读0次

转换onnx后,加载onnx会出现一大串红色warn并且还带着不明字符,需要进行消除。

import onnx

import argparse

def get_args():

    parser = argparse.ArgumentParser()

    parser.add_argument("--input", required=True, help="input model")

    parser.add_argument("--output", required=True, help="output model")

    args = parser.parse_args()

    return args

def remove_initializer_from_input():

    args = get_args()

    model = onnx.load(args.input)

    if model.ir_version < 4:

        print(

            'Model with ir_version below 4 requires to include initilizer in graph input'

        )

        return

    inputs = model.graph.input

    name_to_input = {}

    for input in inputs:

        name_to_input[input.name] = input

    for initializer in model.graph.initializer:

        if initializer.name in name_to_input:

            inputs.remove(name_to_input[initializer.name])

    onnx.save(model, args.output)

if __name__ == '__main__':

    remove_initializer_from_input()

相关文章

网友评论

      本文标题:onnx warn消除

      本文链接:https://www.haomeiwen.com/subject/zyiffdtx.html