在c++中为Eigen实现top_k

Question

我有一个形状为

Eigen::Tensor<float, 3, Eigen::RowMajor, EIGEN_DEFAULT_DENSE_INDEX_TYPE>

 的 3D 特征张量 (

[B, D, C]

)，其中包含代表某些网络输出的置信度得分的浮点数。

我想执行top k并获得最高得分输出的值和索引 沿着第二个维度并将

D

 减少到更小的维度 

k

。所得向量的形状应为 

[B, k, C]

。

在Python（张量流）中它只是（为了简单起见，在本例中为整数）：

data = np.array([[[3, 9, 8, 5],
        [3, 3, 0, 1],
        [5, 5, 1, 9]],

       [[1, 3, 0, 9],
        [8, 3, 7, 7],
        [8, 0, 9, 5]]])


values, indices = tf.math.top_k(tensor, k=2, sorted=False)

结果：

Values

[[[8 9]
  [3 3]
  [5 9]]

 [[3 9]
  [7 8]
  [8 9]]]

Indices

[[[2 1]
  [1 0]
  [0 3]]

 [[1 3]
  [2 0]
  [0 2]]]

提前致谢！

Answer 1

以下代码可以工作，但效率不高：

#include <Eigen/Dense>
#include <algorithm> 
#include <vector>

int B = tensor.dimension(0);
int D = tensor.dimension(1);
int C = tensor.dimension(2);
int k = 2; // Example value for top-k

// Resulting tensors
Eigen::Tensor<float, 3, Eigen::RowMajor> values(B, k, C);
Eigen::Tensor<int, 3, Eigen::RowMajor> indices(B, k, C);

for (int b = 0; b < B; ++b) {
    for (int c = 0; c < C; ++c) {
        // Extract a slice for the current batch element and class
        Eigen::Tensor<float, 1, Eigen::RowMajor> slice = data.chip(b, 0).chip(0, c);

        // Vector to store index-value pairs 
        std::vector<std::pair<float, int>> index_value_pairs;
        index_value_pairs.reserve(D);  // Pre-allocation for efficiency

        // Store index-value pairs
        for (int d = 0; d < D; ++d) {
            index_value_pairs.emplace_back(slice(d), d);
        }

        // Partial sort (only the top-k) based on the values
        std::partial_sort(index_value_pairs.begin(), index_value_pairs.begin() + k, index_value_pairs.end(),
                          [](const std::pair<float, int>& a, const std::pair<float, int>& b) {
                              return a.first > b.first; // Sort by value in descending order
                          });

        // Fill the results 
        for (int i = 0; i < k; ++i) {
            values(b, i, c) = index_value_pairs[i].first;
            indices(b, i, c) = index_value_pairs[i].second;
        }

    }
}