如何在C++ Eigen库中实现张量depth轴最大值索引（类似numpy.argmax()）

嘿，我来帮你搞定这个问题！Eigen的Tensor模块确实没有像numpy那样直接的argmax(axis)语法，但我们可以用它内置的操作或者手动实现来达到同样的效果，下面分两种情况给你说明：

方法一：使用Eigen内置的argmax方法（推荐） #

如果你使用的是Eigen 3.3.7及以上版本，Tensor模块已经提供了argmax方法，可以直接指定要计算的轴。注意要先确认你的张量维度顺序是否正确——你提到张量维度是(行数=200，列数=200，depth=4)，那构造时应该写成Eigen::Tensor<double, 3> table(200, 200, 4);（Eigen的张量维度参数顺序是从最外层到最内层，对应你的行、列、depth）。

完整代码示例：

#include <Eigen/Tensor>
#include <iostream>

int main(){
    // 按照你的需求定义张量：行数200，列数200，depth=4
    Eigen::Tensor<double, 3> table(200, 200, 4);
    table.setRandom();
    
    // 指定要计算最大值索引的轴：depth轴是第三个轴，索引为2（Eigen轴从0开始计数）
    const int target_axis = 2;
    // 计算得到一个2D张量，每个元素对应(行,列)位置上depth轴的最大值索引
    Eigen::Tensor<int, 2> max_indices = table.argmax(target_axis);
    
    // 测试输出：打印(0,0)位置的最大值索引
    std::cout << "Max index at (row=0, col=0): " << max_indices(0, 0) << std::endl;
    
    return 0;
}

注意事项： #

• 如果你的张量实际定义是table(4,200,200)（即depth是第一个轴），那target_axis要改成0，这样才会沿depth轴计算。
• 这个方法是Eigen内部优化过的，效率比手动遍历高很多，适合处理大张量。

方法二：手动实现（兼容旧版Eigen） #

如果你的Eigen版本较低，不支持argmax方法，可以手动遍历每个(行,列)位置，在depth轴上查找最大值的索引：

#include <Eigen/Tensor>
#include <iostream>

int main(){
    Eigen::Tensor<double, 3> table(200, 200, 4);
    table.setRandom();
    
    // 初始化存储索引的2D张量
    Eigen::Tensor<int, 2> max_indices(200, 200);
    
    // 遍历所有行和列
    for(int row = 0; row < 200; ++row){
        for(int col = 0; col < 200; ++col){
            // 获取当前(行,列)对应的depth轴切片
            auto depth_slice = table.chip(row, 0).chip(col, 0);
            
            double max_val = depth_slice(0);
            int max_idx = 0;
            // 遍历depth轴找最大值索引
            for(int d = 0; d < 4; ++d){
                if(depth_slice(d) > max_val){
                    max_val = depth_slice(d);
                    max_idx = d;
                }
            }
            max_indices(row, col) = max_idx;
        }
    }
    
    std::cout << "Max index at (row=0, col=0): " << max_indices(0, 0) << std::endl;
    
    return 0;
}

说明： #

• chip方法用于从张量中提取指定轴的切片，table.chip(row, 0)提取第row行的2D张量，再chip(col, 0)提取第col列的1D张量（也就是depth轴的所有元素）。
• 这种方法虽然代码多一点，但兼容性更好，适合旧版本Eigen。

内容的提问来源于stack exchange，提问作者Pouyan