如何用另一个张量对 PyTorch 张量进行切片？

Question

我有：

inp =  torch.randn(4, 1040, 161)

我还有另一个indices带有值的张量：

tensor([[124, 583, 158, 529],
        [172, 631, 206, 577]], device='cuda:0')

我想要相当于：

inp0 = inp[:,124:172,:]
inp1 = inp[:,583:631,:]
inp2 = inp[:,158:206,:]
inp3 = inp[:,529:577,:]

除非所有内容加在一起，否则 .size 为[4, 48, 161]. 我怎样才能做到这一点？

目前，我的解决方案是一个for循环：

            left_indices = torch.empty(inp.size(0), self.side_length, inp.size(2))
            for batch_index in range(len(inp)):
                print(left_indices_start[batch_index].item())
                left_indices[batch_index] = inp[batch_index, left_indices_start[batch_index].item():left_indices_end[batch_index].item()]

Answer 1

tensor=tensor.cpu()在这里（编辑：在执行以下操作之前，您可能需要使用将张量复制到 cpu ）：

index = tensor([[124, 583, 158, 529],
    [172, 631, 206, 577]], device='cuda:0')
#create a concatenated list of ranges of indices you desire to slice
indexer = np.r_[tuple([np.s_[i:j] for (i,j) in zip(index[0,:],index[1,:])])]
#slice using numpy indexing
sliced_inp = inp[:, indexer, :]

下面是它的工作原理：

np.s_[i:j]i创建一个从 start=到 end=索引的切片对象（只是一个范围）j。

np.r_[i:j, k:m](i,j)创建切片中所有索引的列表(k,m)（您可以传递更多切片以np.r_将它们一次性连接在一起。这是仅连接两个切片的示例。）

因此，indexer通过连接切片列表（每个切片是一个索引范围）来创建 ALL 索引的列表。

更新：如果您需要删除间隔重叠并对间隔进行排序：

indexer = np.unique(indexer)

如果您想删除间隔重叠但不排序并保持原始顺序（以及第一次出现的重叠）

uni = np.unique(indexer, return_index=True)[1]
indexer = [indexer[index] for index in sorted(uni)]