argmax() method in C++ Eigen Library

user_7347196 0
Viewed 2317

I am using Eigen library for matrix/tensor computation where I want to returns the indices of the maximum values along the depth axis. Similar to what numpy.argmax() does in Python.

Tensor dimension is as follows: (rows = 200, columns = 200, depth=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 tensor python c++
3 Answers

Honestly, this is a roundabout not an satisfying answer for the question itself. Nevertheless, the following approach did work for me.

I could not find what I wanted in Eigen library. Instead, I switched to armadillo library which is pretty similar to numpy having a user-friendly API. Moreover, Comparatively speaking, armadillo is more understandable for someone coming form Python or Matlab background.

In armadillo, finding the argmax is as simple as below: (finding argmax of second row and column across all depth axis)

arma::Cube<double> A(200, 200, 4, arma::fill::randu);
uword i = A(arma::span(1),  arma::span(1), arma::span::all).index_max();
user_73471960

Eigen's tensor library has an argmin/argmax member function, which is unfortunately currently not documented on https://eigen.tuxfamily.org/dox/unsupported/eigen_tensors.html.

Eigen's matrix library can mimic the same behavior via visitor overloads of minCoeff/maxCoeff. See: 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);
}

Outputs:

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
user_88766060

This test suite in the Eigen codebase has a few examples for performing reductions across a subset of axes.

Here is an approach that seems to work:

  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 specifies a list of axes that should be reduced (in this case, just the channels). Returned tensor contains both the "linear" indices and maximum values themselves, so to extract actual channel indices, you might need to iterate over the tensor:

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;                
  }                 
}
user_102300530
Related