tensorflow - 将张量从 NHWC 转置到 NCHW 有什么充分的理由吗？

Question

经常在tensorflow代码中看到transpose的实现。我想知道为什么有人会想要将 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。这些线程被组织成 warp(硬件设计)。

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

总结一下:你想让“threadId.x”成为最内循环，你应该组织数据布局，以便它以合并的方式读取它们。理想的数据结构应该可以访问

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

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

在卷积运算(最常用层的一部分)中，您实际上是在每个 channel 中裁剪一个区域，计算与另一个 channel 中的一个区域(来自另一个张量)的点产生。所以需要快速运行的索引是 height-idx 和 width-idx。最后，您将沿 channel 轴添加(如卷积公式所建议的那样)。这也解释了为什么考虑 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 .