python - 具有可变长度数组的索引多维 torch 张量

Question

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

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

有没有办法用索引 i_0, ..., i_k 有效地索引每个 dim d_0, ..., d_k 上的张量？ (k 仅在运行时可用)

结果应该是:

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

目前我正在创建一个切片元组，每个维度一个:

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

但我更喜欢这样的东西:

tensor[:, *idx]

---

示例:

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

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

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

但在一般情况下，我不知道 indexes 有多长。

---

注意: this answer没有帮助，因为它索引所有维度，并且不适用于适当的子集!

Answer 1

不幸的是你can't provide ¹ 索引的切片和迭代器的混合(例如 a[:,*idx])。但是，您可以通过将其包装在方括号中以转换为迭代器来实现几乎相同的效果:

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

---

<支持>

1. >

   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