C++ Eigen 库中的 argmax() 方法

argmax() method in C++ Eigen Library

我正在使用 Eigen 库进行 matrix/tensor 计算,我想 returns 沿深度轴的最大值索引。类似于 numpy.argmax() 在 Python 中的作用。

张量维度如下: (行数 = 200,列数 = 200,深度 = 4)

#include <Eigen/Dense>
int main(){
   Eigen::Tensor<double, 3> table(4,200,200);
   table.setRandom();
   // How can I do this task for axis = 2, i.e depth of a tensor?
   // int max_axis = table.argmax(ax=2); 
   return 0;
}

老实说,这是一个迂回曲折的问题本身,并不是一个令人满意的答案。不过,以下方法对我有用。

我在 Eigen 库中找不到我想要的东西。相反,我切换到 armadillo library which is pretty similar to numpy having a user-friendly API. Moreover, Comparatively speaking, armadillo 对于来自 Python 或 Matlab 背景的人来说更容易理解。

在犰狳中,找到 argmax 很简单,如下所示: (在所有深度轴上查找第二行和第二列的 argmax)

arma::Cube<double> A(200, 200, 4, arma::fill::randu);
uword i = A(arma::span(1),  arma::span(1), arma::span::all).index_max();
Eigen 代码库中的

This test suite 有几个示例,用于跨轴子集执行缩减。

这是一种似乎有效的方法:

  std::array<int, 1> reduce_dims{2};                                            
  Eigen::Tensor<Eigen::Tuple<Eigen::Index, double>, 2> reduced =                 
      table.index_tuples().reduce(reduce_dims,                                  
                                  Eigen::internal::ArgMaxTupleReducer<          
                                      Eigen::Tuple<Eigen::Index, double> >());   

reduce_dims 指定应减少的轴列表(在本例中,仅通道)。 返回的张量包含“线性”索引和最大值本身,因此要提取实际通道索引,您可能需要迭代张量:

for (int c = 0; c < reduced.dimension(1); ++c) {                                 
  for (int r = 0; r < reduced.dimension(0); ++r) {                            
    Eigen::Index argmax_channel_index =                                       
        reduced(r, c).first / (table.dimension(0) * table.dimension(1));      
    std::cout << "argmax channel: " << argmax_channel_index << " "            
              << "max: " << reduced(r, c).second << std::endl;                
  }                 
}

Eigen 的张量库有一个 argmin/argmax 成员函数,不幸的是目前没有在 https://eigen.tuxfamily.org/dox/unsupported/eigen_tensors.html 上记录。

Eigen 的矩阵库可以通过 minCoeff/maxCoeff 的访问者重载来模仿相同的行为。参见:https://eigen.tuxfamily.org/dox/group__TutorialReductionsVisitorsBroadcasting.html

#include <Eigen/Dense>
#include <unsupported/Eigen/CXX11/Tensor>
#include <iostream>

#define STR_(x)  #x
#define STR(x)  STR_(x)
#define PRINT(x)  std::cout << STR(x) << ":\n" << (x) << std::endl

int main()
{
    using namespace Eigen;
    using T = int;
    using S = Sizes<2, 3>;

    S const sizes{};

    T constexpr data[S::total_size]{
        8, 4,
        1, 6,
        9, 2,
    };

    Map<MatrixX<T> const> const matrix(data, sizes[0], sizes[1]);
    PRINT(matrix);

    RowVector2<Index> argmax{};
    matrix.maxCoeff(&argmax.x(), &argmax.y());
    PRINT(argmax);

    VectorX<Index> argmax0{matrix.cols()};
    for (Index col = 0; col < matrix.cols(); ++col)
        matrix.col(col).maxCoeff(&argmax0[col]);
    PRINT(argmax0);

    VectorX<Index> argmax1{matrix.rows()};
    for (Index row = 0; row < matrix.rows(); ++row)
        matrix.row(row).maxCoeff(&argmax1[row]);
    PRINT(argmax1);

    TensorMap<Tensor<T const, S::count>> const tensor(data, sizes);
    PRINT(tensor);
    PRINT(tensor.argmax());
    PRINT(tensor.argmax(0));
    PRINT(tensor.argmax(1));

    // Note that tensor.argmax() is the index for a 1D view of the data:
    Index const matrix_index = sizes.IndexOfColMajor(std::array{argmax.x(), argmax.y()});
    Index const tensor_index = Tensor<Index, 0>{tensor.argmax()}();
    PRINT(matrix_index == tensor_index);
}

输出:

matrix:
8 1 9
4 6 2
argmax:
0 2
argmax0:
0
1
0
argmax1:
2
1
tensor:
8 1 9
4 6 2
tensor.argmax():
4
tensor.argmax(0):
0
1
0
tensor.argmax(1):
2
1
matrix_index == tensor_index:
1