Python脚本
import onnx
from onnx.shape_inference import infer_shapes
import numpy as npfrom google.protobuf.json_format import MessageToJson, Parse
import argparse
import osdef convertToJson(onnx_model_path):onnx_model = onnx.load(onnx_model_path)message = MessageToJson(onnx_model)dir_path, model_name = os.path.split(onnx_model_path)json_model_name = model_name.replace('.onnx', '.json')json_model_path = dir_path + "/" + json_model_namewith open(json_model_path, "w") as fo:fo.write(message)print("Save json to ", json_model_path)def parseArgs():parser = argparse.ArgumentParser(description='ONNX model to json')parser.add_argument("--onnx_model", type=str, required=True, help="Root path of ONNX model.")args = parser.parse_args()print("ONNX model: ", args.onnx_model)return args.onnx_modelif __name__=="__main__":print(">>>>>>>>>>> Begin ONNX model convert <<<<<<<<<<<<<")onnx_model_path = parseArgs()if not os.path.isabs(onnx_model_path):raise ValueError("ERROR! --onnx_model should be the root path.")convertToJson(onnx_model_path)
使用方法
python3 onnx2json --onnx_model /../../test.onnx