是否有任何充分的理由将张量从 NHWC 转置为 NCHW？

Question

我经常在张量流代码中看到转置实现。我想知道为什么有人想要将 NHWC 张量转置为 NCHW。请给我一个好的例子以及背后的原因。

Answer 1

而不是引用文档。您应该了解 CUDA 的工作原理并思考如何实现大多数操作。

NCHW 通常比 NHWC 更快的原因在于 CUDA 内核的编写方式。在 CUDA 中，您需要指定每个线程正在做什么

const int threads = 32;
dim3 block(threads, threads);
dim3 grid(up2(W / 2, threads), up2(H, threads), B);

kernel<Dtype> <<< grid, block>>> (args ...)

这里你得到3个索引threadId.z, threadId.y, threadId.x。这些线程按扭曲组织（硬件设计）。

并且您想要合并内存事务，这意味着线程以这样的方式排序，以便 GPU 可以很好地快速运行。

总结一下：您希望“threadId.x”成为最内循环，并且您应该组织数据布局，以便以合并的方式读取它们。理想的数据结构应该通过以下方式访问

b * C * H * W + c * H * W + h * W + w

其中小写字母表示索引，大写字母表示形状（例如，0 <= w < W）。

在卷积运算（最常用层的一部分）中，您本质上要做的是裁剪每个通道中的一个区域，并与另一个通道中的区域（来自另一个张量）计算点产生。因此需要疯狂快速运行的索引是 height-idx 和 width-idx。最后，您将沿通道轴相加（如卷积公式所示）。这也解释了为什么考虑 NWHC、NCWH 没有什么区别。

这会影响您对数据排序的方式。这就是您想要采用我上面描述的内存布局的原因。

最糟糕的布局是：

H, C, B, in threadId.z, threadId.y, threadId.x

最好的布局是：

B, C, H in threadId.z, threadId.y, threadId.x

对于 GEMM 来说也是如此（大部分）（这里应该转置一个矩阵）。没有可用的 CuDNN 源。但您可能有兴趣研究cutlass。