基本应用
onnx提供的shape_inference的结果经常不完整, 比如这个ssd-12.onnx(models/ssd-12.onnx at main · onnx/models · GitHub):
如果网络中的tensor shape没有推理完全, 那么也没法进行运算量的统计.
刚开源的onnx-tool很好的解决了很多onnx.shape_inference失败的例子. 地址如下:
在ssd这个onnx模型上, onnx-tool可以推理出完整的tensor shapes:
除此之外, 还能够统计出模型每个算子的MACs(浮点乘加数, 和Flops的关系一般是1 MACs=2 Flops)和Params(参数量).
动态输入
上面是基础玩法, 是固定的输入tensor shapes的情况. 如果输入的tensor是dynamic shapes. onnx.shape_inference是不支持你传入动态shape的.
如果像下面这样, 模型(https://github.com/PeterL1n/RobustVideoMatting/releases/download/v1.0.0/rvm_mobilenetv3_fp32.onnx)的输入里面有一个是Resize的scale系数
这个downsample_ratio就是一个动态输入的scale系数. 这种方式onnx.shape_inference就更不支持了.
但是onnx-tool是支持的, 只需要传入一个输入tensor的map就行, 代码也很简单:
自定义Op支持
onnx-tool还可以直接注册一个新的op来支持你自己的网络.
假设你的网络中有一个op是自己定义的, 或者你的onnx版本很新, 有一个onnx-tool还没有支持的op.
比如下面:
这个CropPlugin目前还没有被onnx-tool给实现, 如果直接去shape_infer会报下面的错误:
NotImplementedError: node CropPlugin is not registed for profiling!!! Use NODEPROFILER_REGISTRY to register your profiler for this node.
这时候, 你也不需要等onnx-tool的作者发布下个版本来支持这个op, 你可以直接把这个op在你本地加上就好, 如下:
这样就可以去适配你的ONNX模型的shape inference和MACs统计.
总结
过去一般在进行模型运算量统计的时候, 除了pytorch-opcounter外, 基本没什么好用的工具. onnx-tool算是填补了onnx模型在这一块的空缺, 以及onnx.shape_inference一直没做好的缺陷.
Original: https://blog.csdn.net/luoyu510183/article/details/125561700
Author: Mr_L_Y
Title: ONNX模型tensor shapes inference和Flops统计工具
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/651976/
转载文章受原作者版权保护。转载请注明原作者出处！