c++ - 如何使用 c++ api 获取 rank 3 tensorflow::Tensor 的最大值的索引？

Question

假设我有一个 Tensor t其形状为 {3, 4, 5} .我想找到第一维和第二维的最大值，因此此操作的结果将是形状为 {5, 2} 的矩阵.

到目前为止，我已经尝试通过获取 Eigen::Tensor 来实现这一点来自 tensorflow::Tensor并使用 maxCoeff在一个循环中，所以:

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

但这行不通，因为 t.tensor<float, 3>()返回 Eigen::TensorMap<Eigen::Tensor<float, 3, 1, long>, 16, MakePointer> ，不是 Eigen::Tensor .关于 Eigen::TensorMap 的文档似乎不多类。

我怎样才能得到 Eigen::Tensor来自 Eigen::TensorMap或者做我想用 tensorflow API 做的事情？

Answer 1

为什么不用 axis=0 调用 tf.argmax？