python - pytorch/numpy 中具有任意和可变维数的部分切片

Question

给定 numpy(或 pytorch)中的二维张量，我可以一次沿所有维度进行部分切片，如下所示:

>>> import numpy as np
>>> a = np.arange(2*3).reshape(2,3)
array([[ 0,  1,  2,  3],
       [ 4,  5,  6,  7],
       [ 8,  9, 10, 11]])
>>> a[1:,1:]
array([[ 5,  6,  7],
       [ 9, 10, 11]])

如果我在实现时不知道维数，如何实现相同的切片模式，无论张量中的维数有多少？ (即，如果 a 只有一维，我想要 a[1:]，对于二维，我想要 a[1:,1:]，a[1:,1:,1:] 对于三个维度，依此类推)

如果我可以用如下所示的一行代码来完成它，那就太好了，但这是无效的:

a[(1:,) * len(a.shape)]  # SyntaxError: invalid syntax

我对适用于 pytorch 张量的解决方案特别感兴趣(只需将上面的 numpy 替换为 torch，示例是相同的)，但我认为如果该解决方案同时适用于 numpy 和 pytorch，则可能也是最好的。

Answer 1

答案:制作 slice 的元组对象可以解决这个问题:

a[(slice(1,None),) * len(a.shape)]

说明:slice 是一个内置的 python 类(不依赖于 numpy 或 pytorch)，它提供了用于描述切片的下标表示法的替代方法。 The answer至 a different question建议使用它作为在 python 变量中存储切片信息的方式。 python glossary指出

The bracket (subscript) notation uses slice objects internally.

由于numpy ndarrays的__getitem__方法和 pytorch tensors支持切片的多维索引，它们也必须支持切片对象的多维索引，因此我们可以将这些切片创建一个具有正确长度的元组。

顺便说一句，您可以通过创建一个虚拟类来了解 python 如何使用切片对象，如下所示，然后对其进行切片:

class A(object):
    def __getitem__(self, ix):
        return ix

print(A()[5])  # 5
print(A()[1:])  # slice(1, None, None)
print(A()[1:,1:])  # (slice(1, None, None), slice(1, None, None))
print(A()[1:,slice(1,None)])  #  (slice(1, None, None), slice(1, None, None))