Onnx模型优化

最新推荐文章于 2024-07-24 14:12:43 发布

周星星✘

最新推荐文章于 2024-07-24 14:12:43 发布

阅读量600

点赞数 5

文章标签：人工智能 python onnx

本文链接：https://blog.csdn.net/qq_45152711/article/details/136899154

版权

""" ONNX optimization script

Run ONNX models through the optimizer to prune unneeded nodes, fuse batchnorm layers into conv, etc.

NOTE: This isn't working consistently in recent PyTorch/ONNX combos (ie PyTorch 1.6 and ONNX 1.7),
it seems time to switch to using the onnxruntime online optimizer (can also be saved for offline).

Copyright 2020 Ross Wightman
"""
import argparse
import warnings

import onnx
import onnxoptimizer as optimizer
# from onnx import optimizer
from parse_args import parse_args, getModel

# Load ONNX model
onnx_file_name = "./save_weights/best_model.onnx"
output_name = "./save_weights/optimize.onnx"


def traverse_graph(graph, prefix=''):
    content = []
    indent = prefix + '  '
    graphs =