在 PyTorch 中使用张量索引多维张量 [英] Indexing a multi-dimensional tensor with a tensor in PyTorch

问题描述

我有以下代码:

a = torch.randint(0,10,[3,3,3,3])
b = torch.LongTensor([1,1,1,1])

我有一个多维索引 b 并且想用它来选择 a 中的单个单元格.如果 b 不是张量，我可以这样做:

I have a multi-dimensional index b and want to use it to select a single cell in a. If b wasn't a tensor, I could do:

a[1,1,1,1]

哪个返回正确的单元格，但是:

Which returns the correct cell, but:

a[b]

不起作用，因为它只选择了 a[1] 四次.

Doesn't work, because it just selects a[1] four times.

我该怎么做?谢谢

推荐答案

更优雅(更简单)的解决方案可能是简单地将 b 转换为元组:

A more elegant (and simpler) solution might be to simply cast b as a tuple:

a[tuple(b)]
Out[10]: tensor(5.)

我很想知道这是如何与常规"numpy 一起工作的，并在此处找到了一篇相关文章，对此进行了很好的解释.

I was curious to see how this works with "regular" numpy, and found a related article explaining this quite well here.

这篇关于在 PyTorch 中使用张量索引多维张量的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！