对Resnet50.onnx模型进行BN和卷积层的融合
一、准备工作
安装ONNX
You can then install ONNX from PyPi (Note: Set environment variable ONNX_ML=1
for onnx-ml):
pip install onnx
You can also build and install ONNX locally from source code:
git clone https://github.com/onnx/onnx.git
cd onnx
git submodule update --init --recursive
python setup.py install
二、源码
import onnx import os from onnx import optimizer # Preprocessing: load the model contains two transposes. # model_path = os.path.join('resources', 'two_transposes.onnx') # original_model = onnx.load(model_path) original_model = onnx.load("resnet50.onnx")
print('The model before optimization:\n\n{}'.format(onnx.helper.printable_graph(original_model.graph)))
# A full list of supported optimization passes can be found using get_available_passes() all_passes = optimizer.get_available_passes() print("Available optimization passes:") for p in all_passes: print('\t{}'.format(p)) print() # Pick one pass as example passes = ['fuse_add_bias_into_conv'] # Apply the optimization on the original serialized model optimized_model = optimizer.optimize(original_model, passes) print('The model after optimization:\n\n{}'.format(onnx.helper.printable_graph(optimized_model.graph)))
# save new model
onnx.save(optimized_model, "newResnet50.onnx")