C++ 上二维向量的 Argmax

Question

我正在研究 python/pytorch，我有一个例子

2d vector a
 
    |
    v
  dim-0  ---> -----> dim-1 ------> -----> --------> dim-1
    |   [[-1.7739,  0.8073,  0.0472, -0.4084],
    v    [ 0.6378,  0.6575, -1.2970, -0.0625],
    |    [ 1.7970, -1.3463,  0.9011, -0.8704],
    v    [ 1.5639,  0.7123,  0.0385,  1.8410]]
    |
    v

那么索引为1的argmax就是

# argmax (indices where max values are present) along dimension-1
In [215]: torch.argmax(a, dim=1)
Out[215]: tensor([1, 1, 0, 3])

我的问题是，给定上面的二维向量 a，我如何在 C++ 上实现 argmax 函数来提供与上面相同的输出？谢谢阅读这就是我所做的

   

    vector<vector<float>> a_vect

    {
      {-1.7739,  0.8073,  0.0472, -0.4084},
      {0.6378,  0.6575, -1.2970, -0.0625},
      {1.7970, -1.3463,  0.9011, -0.8704},
      {1.5639,  0.7123,  0.0385,  1.8410}
     };
     std::vector<int>::iterator max = max_element(a_vect.begin() , a_vect.end()-a_vect.begin());

Answer 1

您可以使用std::max_element在每个子向量中找到索引

#include <algorithm>
#include <iostream>
#include <vector>
using std::vector;

int main()
{
  vector<vector<float>> a_vect=
  {
    {-1.7739,  0.8073,  0.0472, -0.4084},
    {0.6378,  0.6575, -1.2970, -0.0625},
    {1.7970, -1.3463,  0.9011, -0.8704},
    {1.5639,  0.7123,  0.0385,  1.8410}
  };
  
  vector<int> max_index;
  for(auto& v:a_vect)
    max_index.push_back(std::max_element(v.begin(),v.end())-v.begin());
  
  for(auto i:max_index)
    std::cout << i << ' '; // 1 1 0 3 
}

C++ 上二维向量的 Argmax

Argmax of 2d vector on C++

c++

pytorch