对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)

猜你喜欢

转载自blog.csdn.net/weixin_45453121/article/details/130923911