读懂NCHW和NHWC

@张小伙

已于 2024-10-08 15:54:11 修改

阅读量833

点赞数 5

文章标签：剪枝

于 2024-10-08 10:25:01 首次发布

本文链接：https://blog.csdn.net/weixin_64043217/article/details/142752955

版权

vscode ssh连接防火墙关闭

sudo ufw allow ssh打开即可

TensorRT推理参考:

下面很好的描述了各种不同的格式。

参考文档

Tensor 内存布局 — MegEngine 1.13.2 文档

对于”NCHW” 而言，其同一个通道的像素值连续排布，更适合那些需要对 每个通道单独做运算 的操作，比如”MaxPooling”。

对于”NHWC”而言，其不同通道中的同一位置元素顺序存储，因此更适合那些需要对 不同通道的同一像素做某种运算 的操作，比如“Conv”。

在GPU上训练时，输入数据格式采用NCHW格式，在推理结果输出时，返回的数据为NHWC格式。

N代表数量， C代表channel，H代表高度，W代表宽度。

NCHW其实代表的是[W H C N]，第一个元素是000，第二个元素是沿着w方向的，即001，这样下去002 003，再接着呢就是沿着H方向，即004 005 006 007...这样到019后，沿C方向，轮到了020，之后021 022 ...一直到319，然后再沿N方向。

NHWC的话以此类推，代表的是[C W H N]，第一个元素是000，第二个沿C方向，即020，040, 060..一直到300，之后沿W方向，001 021 041 061...301..到了303后，沿H方向，即004 024 .。。304.。最后到了319，变成N方向，320,340....

图中所示a就是指四个维度。

那么在不同的硬件加速的情况下，选用的类型不同，在intel GPU加速的情况下，因为GPU对于图像的处理比较多，希望在访问同一个channel的像素是连续的，一般存储选用NCHW，这样在做CNN的时候，在访问内存的时候就是连续的了，比较方便。

NHWC 的访存局部性更好（每三个输入像素即可得到一个输出像素），NCHW 则必须等所有通道输入准备好才能得到最终输出结果，需要占用较大的临时空间。