c++ - 沿多维数组的任意轴减少(求和)

Question

我想沿着可能具有任意维度的多维矩阵的任意轴(例如，10 维数组的第 5 轴)执行总和缩减。矩阵使用行优先格式存储，即作为 vector 以及沿每个轴的步幅。

我知道如何使用嵌套循环执行此缩减(请参见下面的示例)，但这样做会导致硬编码轴(缩减沿下方的轴 1)和任意数量的维度(下方的 4)。 如何在不使用嵌套循环的情况下对此进行概括？

---

#include <iostream>
#include <vector>

int main()
{
  // shape, stride & data of the matrix

  size_t shape  [] = { 2, 3, 4, 5};
  size_t strides[] = {60,20, 5, 1};

  std::vector<double> data(2*3*4*5);

  for ( size_t i = 0 ; i < data.size() ; ++i ) data[i] = 1.;

  // shape, stride & data (zero-initialized) of the reduced matrix

  size_t rshape  [] = { 2, 4, 5};
  size_t rstrides[] = {20, 5, 1};

  std::vector<double> rdata(2*4*5, 0.0);

  // compute reduction

  for ( size_t a = 0 ; a < shape[0] ; ++a )
    for ( size_t c = 0 ; c < shape[2] ; ++c )
      for ( size_t d = 0 ; d < shape[3] ; ++d )
        for ( size_t b = 0 ; b < shape[1] ; ++b )
          rdata[ a*rstrides[0]                 + c*rstrides[1] + d*rstrides[2] ] += \
          data [ a*strides [0] + b*strides [1] + c*strides [2] + d*strides [3] ];

  // print resulting reduced matrix

  for ( size_t a = 0 ; a < rshape[0] ; ++a )
    for ( size_t b = 0 ; b < rshape[1] ; ++b )
      for ( size_t c = 0 ; c < rshape[2] ; ++c )
        std::cout << "(" << a << "," << b << "," << c << ") " << \
        rdata[ a*rstrides[0] + b*rstrides[1] + c*rstrides[2] ] << std::endl;

  return 0;
}

---

注意:我想避免“解压缩”和“压缩”计数器。我的意思是我可以用伪代码来做:

for ( size_t i = 0 ; i < data.size() ; ++i ) 
{
  i -> {a,b,c,d}

  discard "b" (axis 1) -> {a,c,d}

  rdata(a,c,d) += data(a,b,c,d)
}

Answer 1

我不知道这段代码的效率如何，但在我看来，它肯定是精确的。

发生了什么事？

关于 adjusted_strides 的一点点:

对于 axis_count = 4 , adjusted_strides尺寸 5 ，其中:

 adjusted_strides[0] = shape[0]*shape[1]*shape[2]*shape[3];
 adjusted_strides[1] = shape[1]*shape[2]*shape[3];
 adjusted_strides[2] = shape[2]*shape[3];
 adjusted_strides[3] = shape[3];
 adjusted_strides[4] = 1;

让我们以维数为4的例子为例多维数组 ( A ) 的形状是 n0, n1, n2, n3 .

当我们需要将此数组转换为形状为 B 的另一个多维数组 ( n0, n2, n3) 时(压缩 axis = 1 (0-based) )，那么，我们尝试如下进行:

对于 A 的每个索引我们试图找到它在 B 中的位置.让A[i][j][k][l]是 A 中的任何元素.它在 flat_A 中的位置将是 A[i*n1*n2*n3 + j*n2*n3 + k*n3 + l]

idx = i*n1*n2*n3 + j*n2*n3 + k*n3 + l;

在压缩数组中 B ，此元素将成为(或添加到)B[i][k][l] 的一部分.在 flat_B索引是 new_idx = i*n2*n3 + k*n3 + l; .

我们如何形成new_idx来自 idx ?

1. 压缩轴之前的所有轴都具有压缩轴的形状作为其乘积的一部分。在我们的示例中，我们必须删除轴 1 , 所以所有在第一个轴之前的轴(这里只有一个: 0th axis )由 i 表示), 有 n1作为产品的一部分 ( i*n1*n2*n3 )。

2. 压缩轴之后的所有轴不受影响。

3. 最后，我们需要做两件事:

   1. 在要压缩的轴的索引之前隔离轴的索引，并删除该轴的形状:

      整数除法:idx / (n1*n2*n3); (== idx / adjusted_strides[1])。

      我们只剩下 i ，可以根据新形状重新调整(乘以 n2*n3 ):我们得到

      i*n2*n3 (== i * adjusted_strides[2])。

   2. 我们在压缩轴之后隔离轴，这些轴不受其形状的影响。

      idx % (n2*n3) (== idx % adjusted_strides[2])

      这给了我们 k*n3 + l .

   3. 将步骤i. 和ii. 的结果相加得到:

      computed_idx = i*n2*n3 + k*n3 + l;

      这与 new_idx 相同.所以，我们的转换是正确的:)。

代码:

备注:ni指new_idx .

  size_t cmp_axis = 1, axis_count = sizeof shape/ sizeof *shape;
  std::vector<size_t> adjusted_strides;
  //adjusted strides is basically same as strides
  //only difference being that the first element is the 
  //total number of elements in the n dim array.

  //The only reason to introduce this array was
  //so that I don't have to write any if-elses
  adjusted_strides.push_back(shape[0]*strides[0]);
  adjusted_strides.insert(adjusted_strides.end(), strides, strides + axis_count);
  for(size_t i = 0; i < data.size(); ++i) {
    size_t ni = i/adjusted_strides[cmp_axis]*adjusted_strides[cmp_axis+1] + i%adjusted_strides[cmp_axis+1];
    rdata[ni] += data[i];
  }

输出(轴=1)

(0,0,0) 3
(0,0,1) 3
(0,0,2) 3
(0,0,3) 3
(0,0,4) 3
(0,1,0) 3
(0,1,1) 3
(0,1,2) 3
(0,1,3) 3
(0,1,4) 3
(0,2,0) 3
(0,2,1) 3
(0,2,2) 3
(0,2,3) 3
(0,2,4) 3
(0,3,0) 3
(0,3,1) 3
(0,3,2) 3
...

已测试 here .

如需进一步阅读，请参阅 this .