由另一个多维张量索引多维火炬张量 [英] Index multidimensional torch tensor by another multidimensional tensor

问题描述

我在 pytorch 中有一个张量 x 比方说形状 (5,3,2,6) 和另一个形状张量 idx (5,3,2,1) 其中包含第一个张量中每个元素的索引.我想用第二个张量的索引对第一个张量进行切片.我试过 x= x[idx] 但是当我真的希望它的形状为 (5,3,2) 或 (5,3,2,1) 时，我得到了一个奇怪的维度.

I have a tensor x in pytorch let's say of shape (5,3,2,6) and another tensor idx of shape (5,3,2,1) which contain indices for every element in first tensor. I want a slicing of the first tensor with the indices of the second tensor. I tried x= x[idx] but I get a weird dimensionality when I really want it to be of shape (5,3,2) or (5,3,2,1).

我会尝试举一个更简单的例子:说吧

I'll try to give an easier example: Let's say

x=torch.Tensor([[10,20,30],
                 [8,4,43]])
idx = torch.Tensor([[0],
                    [2]])

我想要类似的东西

y = x[idx]

这样 'y' 输出 [[10],[43]] 或类似的东西.

such that 'y' outputs [[10],[43]] or something like.

索引表示最后一维所需元素的位置.对于上面的示例，其中 x.shape = (2,3) 最后一个维度是列，然后 'idx' 中的索引是列.我想要这个，但超过 2 个维度

The indices represent the position of the wanted elements the last dimension. for the example above where x.shape = (2,3) the last dimension are the columns, then the indices in 'idx' is the column. I want this but for more than 2 dimensions

推荐答案

根据我从评论中的理解，您需要 idx 在最后一个维度中作为索引，并且每个索引在 idx 中 对应于 x 中的相似索引(最后一个维度除外).在这种情况下(这是 numpy 版本，您可以将其转换为 Torch):

From what I understand from the comments, you need idx to be index in the last dimension and each index in idx corresponds to similar index in x (except for the last dimension). In that case (this is the numpy version, you can convert it to torch):

ind = np.indices(idx.shape)
ind[-1] = idx
x[tuple(ind)]

输出:

[[10]
 [43]]

这篇关于由另一个多维张量索引多维火炬张量的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！