具有可变长度数组的索引多维火炬张量 [英] Index multidimensional torch tensor with array of variable length

问题描述

我有一个索引列表和一个具有张量的张量:

I have a list of indices and a tensor with shape:

shape = [batch_size, d_0, d_1, ..., d_k]
idx = [i_0, i_1, ..., i_k]

有没有一种方法可以用索引 i_0，...，i_k 有效地索引每个暗的 d_0，...，d_k 上的张量?( k 仅在运行时可用)

Is there a way to efficiently index the tensor on each dim d_0, ..., d_k with the indices i_0, ..., i_k? (k is available only at run time)

结果应为:

tensor[:, i_0, i_1, ..., i_k] #tensor.shape = [batch_size]

此刻，我正在创建一个切片的元组，每个维度一个:

At the moment I'm creating a tuple of slices, one for each dimension:

idx = (slice(tensor.shape[0]),) + tuple(slice(i, i+1) for i in idx)
tensor[idx]

但是我更喜欢这样的东西:

but I would prefer something like:

tensor[:, *idx]

---

示例:

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

我只想索引最后一个len(indexes)维，例如:

I would like to index only the last len(indexes) dimensions like:

a[:, indexes[0], indexes[1], indexes[2]]

但是在一般情况下，我不知道索引有多长时间.

but in the general case where I don't know how long indexes is.

注意:此答案没有帮助，因为它索引了全部尺寸，不适用于适当的子集！

Note: this answer does not help since it indexes all the dimensions, and does not work for a proper subset!

推荐答案

很遗憾，您无法提供 ¹ 切片和迭代器的混合以建立索引(例如 a [:，* idx] ).但是，您可以通过将其包装在方括号中以投射到迭代器上来实现几乎相同的目的:

Unfortunately you can't provide¹ a mix of slices and iterators to an indexing (e.g. a[:,*idx]). However, you can achieve almost the same thing by wrapping it in brackets to cast to an iterator:

a[(slice(None), *idx)]

---

在Python中， x [(exp1，exp2，...，expN)] 等效于 x [exp1，exp2，...，expN] ;后者只是前者的语法糖.

In Python, x[(exp1, exp2, ..., expN)] is equivalent to x[exp1, exp2, ..., expN]; the latter is just syntactic sugar for the former.

https://numpy.org/doc/stable/reference/arrays.indexing.html

这篇关于具有可变长度数组的索引多维火炬张量的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！