多维(NCHW)Tensor中的数据摆放顺序变换

遇到一个问题,将n*c*h*w的tensor转换为n*(h*w)*1*c.
tensor储存在一维数组a中,易得坐标(nn,cc,hh,ww)的元素在a中的下标为:

nn*c*h*w + cc*h*w + hh*w + ww

n*(h*w)1*c可以看成n(h*w)*c,四维张量简化为三维,因为h*w原本为两维,压缩成了(h*w)一维.

因此原问题简化为,将n*(h*w)c转换为n(h*w)*c,即交换低两维.

已知高维和低维的区别在于计算下标时的先后顺序,交换维度意味着原数据在一维数组a中的摆放方式也要随之改变,画图说明:

这么一来就非常明确了,用三个嵌套for循环即可完成变换:

//a为原数组,b为变换数组
for (int nn = 0; nn < n; nn++) {
  for (int hw = 0; hw < h*w; hw++) {
    for (int cc = 0; cc < c; cc++) {
      // b下标按n*(h*w)*c顺序
      // a下标按n*c*(h*w)顺序
      b[nn*h*w*c + hw*c + cc] = a[nn*c*h*w + cc*h*w + hw];
    }
  }
}