LibTorch中的tensor.index_select()方法与PyTorch中的用法类似，作用都是在指定的tensor维度dim上按照index值索引向量。先看一下在LibTorch中的声明：
inline Tensor Tensor::index_select(int64_t dim, const Tensor & index)
主要是两个参数，一个是要选择的维度，对于二维tensor来说，0代表按行索引，1代表按列索引，参数index代表索引值，这个参数要重点注意一下！首先index它本身就是一个tensor！另外，还得是int64（kLong）类型的tensor！
下面举例说明：
如果我们要获取一个[shape为(3, 2)]的二维tensor的第列维度上(dim=1)的数据(index=0，即第0列)[shape为(3, 1)]并打印，如果像下面这样写：（from_blob的作用是将数组转化为tensor）
#include "iostream"
#include "torch/script.h"
using namespace torch::indexing;
int main()
{
float a[3][2] = {1};
at::Tensor b = at::from_blob(a, {3, 2}, at::kFloat).index_select(1, at::tensor(0));
std::cout << b << std::endl;
return 0;
}
如果就这么运行，那么bug就会出现了。。。哈哈哈，先看看bug说的啥吧：
terminate called after throwing an instance of 'c10::Error'
what(): index_select(): Expected dtype int64 for index (index_select_out_cpu_ at ../aten/src/ATen/native/TensorAdvancedIndexing.cpp:387)
发现问题没？上面不是说了嘛，index它本身就是一个tensor！另外，还得是int64（kLong）类型的tensor！所以我们还要对at::tensor(0)做一下数据类型的转换，如下面这段代码就正常了：
#include "iostream"
#include "torch/script.h"
using namespace torch::indexing;
int main()
{
float a[3][2] = {{1,2},{3,4},{5,6}};
at::Tensor b = at::from_blob(a, {3, 2}, at::kFloat).index_select(1, at::tensor(0).toType(at::kLong));
std::cout << b << std::endl;
return 0;
}
/*****************运行结果*******************/
1
3
5
[ CPUFloatType{3,1} ]
本文内容由网友自发贡献，版权归原作者所有，本站不承担相应法律责任。如您发现有涉嫌抄袭侵权的内容，请联系:hwhale#tublm.com(使用前将#替换为@)
