Trim the redundant node "identity" in the onnx file

Trim the redundant node "identity" in the onnx file

In the process of model deployment, I found that the converted onnx model contains a large number of "identity" nodes. During the actual measurement, it was found that these nodes will not affect the inference results of onnx, nor will they affect the inference results of the tensorrt inference engine converted from onnx, but will seriously affect the inference speed of onnx. At the same time, the converted tensorrt file is also much larger than normal, and the inference speed will also be seriously affected.

So record the pruning process and code for the "identity" node here.

identity node

I haven't figured out yet why there are a lot of "identity" nodes, but it doesn't seem to affect my final result, just the speed. My onnx visualization looks like this:

After cutting, it looks like this:

After cutting, convert the tensorrt file. The file size is much smaller than before, and the speed is faster.

trim code

I use onnx-simplifier for pruning, first install:

pip install onnx-simplifier

Then you can use it directly:

import onnx
import onnxoptimizer

#加载修剪之前的模型
onnx_model = onnx.load(r"mmdeploy/s2anet2onnx2/resnet50+cat+fam+deform-1024-prune.onnx")

#去掉identity层
all_passes = onnxoptimizer.get_available_passes()
print("Available optimization passes:")
for p in all_passes:
    print('\t{}'.format(p))
print()

#保存修剪后的模型
onnx_optimized = 'mmdeploy/s2anet2onnx2/resnet50+cat+fam+deform-1024-prune.onnx'
passes = ['eliminate_identity']
optimized_model = onnxoptimizer.optimize(onnx_model, passes)
onnx.save(optimized_model, onnx_optimized)