tensorflow tf.space_to_batchnd函数理解

最近需要将tensorflow转ncnn模型, 发现无论是普通的卷积还是分组卷积, 只要有dilation > 1的情况出现, tensorflow在做空洞卷积的时候, 就会执行:

1. SpaceToBatchND操作
2. conv操作
3. BathToSpaceND操作

其中, 第二步的conv操作已经是普通的卷积操作了, 即dilation=1的普通卷积.
tensorflow这样做的目的, 是不改变kernel的值, 而是事先根据dilation的值改变featuremap的值(通过SpaceToBatchND函数), 然后再做普通的卷积, 最后再通过BatchToSpaceND, 将featuremap转换为原有的featuremap直接执行空洞卷积应该得到的结果. 不清楚的看下图:

在featuremap为4x4, kernel 为2x2, dilation=2的情况下, 卷积核在featuremap上的每次取值为上图中包含相同图标的四个值. 如果直接按照这个方式来做空洞卷积, 那么速度就会非常慢, 所以tensorflow推出SpaceToBatchND函数, 根据dilation的值,提前把featuremap上应该参与卷积的位置给提出来, 然后组成一个大的batch, 接着和kernel做普通的dilation=1的卷积, 最后再根据BatchToSpaceND返回原来应有的结果. 理解了这个, 那么这两个函数的具体操作就很容易理解了, 根据官方源码的介绍:
SpaceToBatchND