Search code examples
c++tensorfloweigen

How can I get the indices of the max values of a rank 3 tensorflow::Tensor using the c++ api?

Let's say I have a Tensor t which has shape {3, 4, 5}. I would like to find the max value of the first and second dimensions, so the result of this operation would be a matrix of shape {5, 2}.

So far I have tried to this by getting an Eigen::Tensor from the tensorflow::Tensor and using maxCoeff in a loop, so:

auto t_mapped = t.tensor<float, 3>();
Eigen::Matrix<float, 5, 2> maximums;

for (int i = 0; i < 5; i++){
    MatrixXf::Index maxRow, maxCol;
    t_mapped.maxCoeff(&maxRow, &maxCol);
    maximums(i, 0) = maxRow;
    maximums(i, 1) = maxCol;
}

But this doesn't work because t.tensor<float, 3>() returns an Eigen::TensorMap<Eigen::Tensor<float, 3, 1, long>, 16, MakePointer>, not an Eigen::Tensor. There doesn't appear to be much documentation on the Eigen::TensorMap class.

How can I either get an Eigen::Tensor out of the Eigen::TensorMap or do what I'm trying to do with the tensorflow API?

Solution

  • Why not call tf.argmax with axis=0?