Conv + BN + Act算子融合

首发:知乎
作者:张新栋

我们在算法模型中往往采用了大量的BN层,但是在inference阶段,很多BN层的使用是可以合并到一些线性运算中的(比如文章题图中的Conv + BN + Act组合)。我们将Conv + BN + Act的算子进行融合后,只需要采用一个Conv + inplace activation就可以完成原有的3次Op计算,绝大多数情况下还能得到一致的数值结果(算子融合极少数情况下会出现数值精度异常)。该融合理论简单,融合op不仅减少了BN层的数值操作,还能减少2次访存,美滋滋!理论部分就比较简单啦,可以参考下图:

conv + BN都是线性操作,参数直接一算就融合起来啦。很多框架和开源工作都提供了fuse BN的操作,我们这里和大家讨论一下对tensorflow pb如何进行fuse BN的操作(onnx的话,强烈推荐大老师的onnx-simplifier)。

tensorflow提供了optimize for inference的python接口,我们可以利用这个api进行graph的optimize,其中就有fuseBN的操作。

我们来看看optimize之前graph的拓扑图:

经过tf的optimize之后,graph拓扑图变成下图样子:

大家可以看到,BN层的参数以及运算都融合进Conv2D op和BiasAdd op里啦。整个图瞬间清爽许多。

本文选用一个简单fuseBN的问题,跟大家简单介绍了理论和基于tensorflow pb的例子。大家如果有感兴趣的任务或者话题,也可以下方留言。欢迎大家留言、关注专栏,谢谢大家!

  • 参考

1、https://github.com/daquexian/onnx-simplifier

2、https://github.com/lutzroeder/netron

推荐阅读

  • 基于onnx的网络裁剪
  • RK1808-AI开发手记(二)人脸姿态估计porting(python)
  • RK1808-AI开发手记(一)开篇

专注嵌入式端的AI算法实现,欢迎关注作者微信公众号和知乎嵌入式AI算法实现专栏

更多嵌入式AI相关的技术文章请关注极术嵌入式AI专栏

发表评论

邮箱地址不会被公开。 必填项已用*标注