基本应用

onnx提供的shape_inference的结果经常不完整, 比如这个ssd-12.onnx(models/ssd-12.onnx at main · onnx/models · GitHub):

如果网络中的tensor shape没有推理完全, 那么也没法进行运算量的统计.

刚开源的onnx-tool很好的解决了很多onnx.shape_inference失败的例子. 地址如下:

onnx-tool · PyPI  

在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就行, 代码也很简单:

 如上图中的代码, downsample_ratio是0.25, 也就是输入1080x1920下输出的270x480, 如下图:

 

---

自定义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一直没做好的缺陷.