任意 Dims 的动态切片张量答案

【问题标题】：Dynamically Slice Tensor of Arbitrary Dims任意 Dims 的动态切片张量
【发布时间】：2022-01-25 09:42:03
【问题描述】：

我正在编写一个定义，该定义需要对具有任意维数的张量切片。切片将始终位于 0 的批次维度上。

这是我想要的一个简单示例：

def masktensor(X, array_of_indices):
    return X[array_of_indices, *] # edit * to allow variable number of dims

我希望能够为其提供各种大小的张量，并将索引切片作为一个批次取出。如：

A = torch.rand(1000, 3, 32, 32)
B = torch.rand(1000, 5, 20)

indices = np.arange(10)

A_batch = masktensor(A, indices)
B_batch = masktensor(B, indices)

提前致谢！

【问题讨论】：

标签： python numpy pytorch

【解决方案1】：

经过一番反复试验，我发现... 可以用作切片通配符：

def masktensor(X, array_of_indices):
    return X[array_of_indices, ...] # <---- '...' allows variable dim size

A = torch.rand(1000, 3, 32, 32)
B = torch.rand(1000, 5, 20)

indices = np.arange(10)

print(masktensor(A, indices).size())
print(masktensor(B, indices).size())

产量：

>>> torch.Size([10, 3, 32, 32])
>>> torch.Size([10, 5, 20])

【讨论】：