C++中如何高效实现NCHW格式到NHWC格式的转换

前置说明 #

OpenCV默认读取的图像为HWC布局，若你拿到的是NCHW格式的批量数据，可选择以下几种优化方案替代手写三层循环，性能提升3~20倍不等。

方案1：使用OpenCV DNN模块内置permute接口（最推荐） #

OpenCV官方的cv::dnn::permute接口已经做了多硬件后端（CPU、GPU、NPU）的底层优化，自动处理内存对齐、向量加速、多线程调度，实现成本最低：

#include <opencv2/dnn.hpp>

cv::Mat nchw_blob; // 输入NCHW格式数据，维度[N, C, H, W]
cv::Mat nhwc_blob;
// 新维度索引顺序：原维度0(N)、2(H)、3(W)、1(C)，对应NHWC布局
cv::dnn::permute(nchw_blob, nhwc_blob, cv::Vec<int,4>{0, 2, 3, 1});

若数据存储在GPU显存上，可直接使用cv::cuda::transpose或CUDA官方的张量转置内核，避免CPU-GPU数据拷贝开销。

方案2：Eigen张量重排（适合已有Eigen依赖的项目） #

如果项目已经集成Eigen库，可以直接使用张量的shuffle方法，底层也做了SIMD优化：

#include <Eigen/CXX11/Tensor>

Eigen::Tensor<float, 4, Eigen::RowMajor> nchw_tensor(N, C, H, W);
// 填充输入数据到nchw_tensor
Eigen::array<int, 4> permute_order = {0, 2, 3, 1};
Eigen::Tensor<float, 4, Eigen::RowMajor> nhwc_tensor = nchw_tensor.shuffle(permute_order);

方案3：手动SIMD优化实现（无第三方依赖场景） #

如果不想引入额外依赖，可以手写SIMD向量化实现，x86平台下可使用AVX2指令集，ARM平台下可替换为NEON指令：

• 编译时需要开启对应指令集选项：GCC加-mavx2，MSVC加/arch:AVX2

#include <immintrin.h>

// 示例为float类型的批量数据转置，可根据数据类型、通道数自行扩展
void nchw2nhwc_avx2(float* dst, const float* src, int N, int C, int H, int W) {
    int hw_size = H * W;
    for (int n = 0; n < N; n++) {
        // 每次处理8个像素，对应AVX2 256位寄存器的8个float位宽
        for (int i = 0; i < hw_size; i += 8) {
            for (int c = 0; c < C; c++) {
                __m256 ch_vals = _mm256_loadu_ps(src + n * C * hw_size + c * hw_size + i);
                for (int k = 0; k < 8; k++) {
                    dst[n * hw_size * C + (i + k) * C + c] = ch_vals[k];
                }
            }
        }
    }
}

内容的提问来源于stack exchange，提问作者Moghaddam