NCHW与NC4HW4数据排布在卷积优化上的优劣分析

NCHW与NC4HW4数据排布在卷积优化上的优劣分析


数据布局对性能影响巨大。

先来看一看在 NCHW 的布局下,怎么利用 SIMD 加速 3x3 的 depth-wise 卷积。

首先,读取数据时,需要一次性读取四个 float 作为第一行的数据,后两行的读取也是相似的;此时,读取出的三行数据已经足够计算两列输出,即,可以复用部分数据;而后,为了提高数据复用,会再读取出第四行数据,一次计算两行两列,即,可以引入循环展开;然而,残留的 5~25 和 21~25 亮度眼边界无法利用 SIMD 计算,只能逐一循环读写完成计算;按照这样的方式,就可以相应完成后几个通道的计算。

但是, NCHW 布局下,无法充分利用 SIMD 进行加速,同时,实现优化分支越多,占用包大小也就越多。

NCHW与NC4HW4数据排布在卷积优化上的优劣分析


再来看一看 NC/4HW4 布局下,利用 SIMD 加速的情况又是怎样的。

这里的 "C/4" 指的是按照 4 个通道对齐的方式重排数据。重排所有输入和权重数据后,每次 SIMD 读写都天然是 4 个通道的输入数据和 4 个通道的权重数据。这样,不论 kernel、stride、dilation 怎么变化,我们都可以简单地使用 for 循环和 SIMD 的一套通用优化完成卷积计算。既不会有边缘数据无法加速的问题,也不会对包大小造成影响。