NHWC和NCHW数据排布及转换(模型部署)

news2024/10/5 14:07:22

1.概念

首先这是两种批量图片的数据存储方式，定义了一批图片在计算机存储空间内的数据存储layout。N表示这批图片的数量，C表示每张图片所包含的通道数，H表示这批图片的像素高度，W表示这批图片的像素宽度。其中C表示的通道数可能有多种情况，例如，RGB图片格式的通道为3通道，R（红色）、G（绿色）、B（蓝色）各占一个通道，表示图片中每个像素点都有三个通道值，每个通道值范围是[0~255]，三个通道的叠加呈现出一个像素的颜色。RGB图像还有四通道的表示，除了RGB三通道之外，还有一个alpha通道，表示透明度。如果是灰度图，则只有一个通道。YUV图片也包含了三个通道，与RGB不同的是YUV的数据表现方式有多样组合。

2. 数据layout理解

以RGB数据格式的图像为例，一张RGB图片可以由H、W、C三个参数来描述，图像的高度和宽度分别对应H和W这两个维度，而图像中每一个像素点都有3个值来表示，便是表示红色的R、绿色的G和蓝色的B，每个值的取值范围是[0~255]，如下图所示：

在这里插入图片描述

2. 1 NCHW和NHWC在模型推理中的应用

在实际部署时，比如采用部署框架如NCNN、TensorRT、Caffe等，模型推理后会返回类型基本都是float* 的结果数据data,即存储为一维数据。此时如果知道返回data的数据排布是NCHW还是NHWC,就很容易对结果进行解析。两种格式的数据的排布如下:

在这里插入图片描述

(1) NCHW排布

NCHW类似于python 中numpy.array 4维数组的排布：最内层为H,W,外层是C,最外层是N 。由于返回的data是float*类型，因此将4维数据从内层向外层逐渐拆分出来（先排HW数据），排布为一维数组。
即[R11,R12,R13....R21,R22,R23,R24...G11,G12,G13....G21,G22,G23,G24...B11,B12,B13....B21,B22,B23,B24...] , 每个通道的元素会紧挨在一起。假设RGB图像的大小为HxWx3, 对于NCHW排布的数据，[0:HxW]存储的都是R通道的像素值[HxW:2*HxW]存储的都是G通道的像素值，[HxW*2:3*HxW]存储的是B通道的像素值

在这里插入图片描述

(1) NHWC排布

NHWC类似于python 中numpy.array 4维数组的排布：最内层为C,W最外层是H,N。由于返回的data是float*类型，因此将4维数据从内层向外层逐渐拆分出来(先排C数据)，排布为一维数组。

即[R11,G11,B11,R12,G12,B12,R13,G13,B13...R21,G21,B21,R22,G22,B22....Rij,Gij,Bij], 像素点ij位置上多个通道中的数值连续存储。

在这里插入图片描述
最后是多张图片的存储N，表示一共有N张图片。NHWC和NCHW表示两种图片数据存储方式，应用于不同的硬件加速场景之下。在intel GPU加速的情况下，希望在访问同一个channel的像素是连续的，一般存储选用NCHW也就是输入数据格式为NCHW，这样在做CNN的时候，在访问内存的时候就是连续的了，比较方便。

最佳实践: 设计网络时充分考虑两种格式，最好能灵活切换。在GPU上训练时，输入数据格式采用NCHW格式，在推理结果输出时，返回的数据为NHWC格式。

3. NHWC与NCHW之间转化

两种存储方式展示了图片数据在储存器中的存储方式，这两种存储方式之间可以相互转换，以NHWC转NCHW为例，可以做如下转换：

3.1 NHWC 转 NCHW

int nhwc_to_nchw(float *out_data, float *in_data, int img_h, int img_w) {
  auto *hw1 = out_data;
  auto *hw2 = hw1 + img_h * img_w;
  auto *hw3 = hw2 + img_h * img_w;
  for (int hh = 0; hh < img_h; ++hh) {
    for (int ww = 0; ww < img_w; ++ww) {
      *hw1++ = *(in_data++);    // B
      *hw2++ = *(in_data++);    // G
      *hw3++ = *(in_data++);    // R
    }
  }
  return 0;
}

in_data 为输入的NHWC 排布的数据, out_data为转换后的NCHW数据

3.2 NCHW 转 NHWC

int nchw_to_nhwc(float* out_data, float* in_data, int img_h, int img_w) {
  auto *res = out_data;

  for (int i = 0; i < img_h * img_w*3 ;) {
      res[i] = *(in_data);    
      res[i+1] = *(in_data + img_h * img_w); 
      res[i+2] = *(in_data + 2*img_h * img_w); 
      i +=3;
    }
  }
  return 0;
}