对onnx文件里的冗余节点“identity”进行修剪

对onnx文件里的冗余节点“identity”进行修剪

在进行模型部署的过程中，我发现转换好的onnx模型含有大量的“identity”节点。实测过程中发现，这些节点不会影响onnx的推理结果，也不会影响onnx转换成的tensorrt推理引擎的推理结果，但是会严重影响onnx的推理速度。同时转换成的tensorrt文件也比正常的大很多，推理速度也会受到严重影响。

所以在此记录对“identity”节点的修剪过程和代码。

identity节点

目前我还没搞清楚为什么会出现大量“identity”节点，但是它貌似没有影响我最终的结果，只是影响了速度。我的onnx可视化效果长这样：

剪完之后长这样：

剪完之后再转换tensorrt文件，文件大小要比之前小很多，速度也变快了

修剪代码

我使用的onnx-simplifier进行修剪的，先安装：

pip install onnx-simplifier

然后就可以直接用了：

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)