跨特定维度的显式切片

Question

我有一个3D张量x（例如4x4x100）。我想通过明确选择最后一个维度中的元素来获取其中的一个子集。如果我在最后一个维度上选择相同的元素（例如x[:,:,30:50]，但是我想使用2D张量indices来指定跨越第三维的idx，那么这将很容易。）有一种简单的方法可以在numpy的？

一个更简单的2D示例：

x = [[1,2,3,4,5,6],[10,20,30,40,50,60]]
indices = [1,3]

假设我想从x指定的点开始，在indices的第三维上抓取两个元素。所以我想要的输出是：

[[2,3],[40,50]]

更新：我想我可以使用take()和ravel_multi_index()的组合，但是一些受numpy启发的平台（如PyTorch）似乎没有ravel_multi_index所以我正在寻找替代解决方案

Answer 1

如果“行”的数量不是太大（并且尺寸的大小相对较大），则迭代idx并收集切片并不是一个坏选择。

In [55]: x = np.array([[1,2,3,4,5,6],[10,20,30,40,50,60]])                      
In [56]: idx = [1,3]                                                            
In [57]: np.array([x[j,i:i+2] for j,i in enumerate(idx)])                       
Out[57]: 
array([[ 2,  3],
       [40, 50]])

像这样加入切片只有在大小相同的情况下才有效。

另一种方法是将索引收集到一个数组中，然后进行一次索引。

例如，使用类似的迭代：

idxs = np.array([np.arange(i,i+2) for i in idx])

但广播的增加可能会更好：

In [58]: idxs = np.array(idx)[:,None]+np.arange(2)                              
In [59]: idxs                                                                   
Out[59]: 
array([[1, 2],
       [3, 4]])
In [60]: x[np.arange(2)[:,None], idxs]                                          
Out[60]: 
array([[ 2,  3],
       [40, 50]])

---

ravel_multi_index不难复制（如果你不需要剪裁等）：

In [65]: np.ravel_multi_index((np.arange(2)[:,None],idxs),x.shape)              
Out[65]: 
array([[ 1,  2],
       [ 9, 10]])
In [66]: x.flat[_]                                                              
Out[66]: 
array([[ 2,  3],
       [40, 50]])
In [67]: np.arange(2)[:,None]*x.shape[1]+idxs                                   
Out[67]: 
array([[ 1,  2],
       [ 9, 10]])

Answer 2

沿3D轴：

x = [x[:,i].narrow(2,index,2) for i,index in enumerate(indices)]
x = torch.stack(x,dim=1)

通过枚举，您可以从一个开始切片的位置获取轴和索引的索引。

narrow从一个特定轴的起始索引length给你一个零拷贝的start长切片

你说你想要的：

dim = 2
start = index
length = 2

那么你只需将这些张量叠加回一个3D。

对于pytorch来说，这是我能想到的最少工作密集的事情。

编辑

如果您只是想要沿不同轴的不同索引而且indices是2D张量，您可以这样做：

x = [x[:,i,index] for i,index in enumerate(indices)]
x = torch.stack(x,dim=1)

你真的应该给出一个正确的工作示例，使其不必要地混淆。

Answer 3

下面是如何在numpy中做到这一点，现在关于火炬的线索。

下面从点idx开始沿第三维选取一段长度为n的长度，具体取决于其他两个维度：

# example
a = np.arange(60).reshape(2, 3, 10)
idx = [(1,2,3),(4,3,2)]
n = 4

# build auxiliary 4D array where the last two dimensions represent
# a sliding n-window of the original last dimension
j,k,l = a.shape
s,t,u = a.strides
aux = np.lib.stride_tricks.as_strided(a, (j,k,l-n+1,n), (s,t,u,u))

# pick desired offsets from sliding windows
aux[(*np.ogrid[:j, :k], idx)]
# array([[[ 1,  2,  3,  4],
#         [12, 13, 14, 15],
#         [23, 24, 25, 26]],

#        [[34, 35, 36, 37],
#         [43, 44, 45, 46],
#         [52, 53, 54, 55]]])

Answer 4

我在下面用广播提出：

x = np.array([[1,2,3,4,5,6,7,8,9,10],[10,20,30,40,50,60,70,80,90,100]]) 
i = np.array([1,5])
N = 2 # number of elements I want to extract along each dimension. Starting points specified in i

r = np.arange(x.shape[-1])
r = np.broadcast_to(r, x.shape)

ii = i[:, np.newaxis]
ii = np.broadcast_to(ii, x.shape)

mask = np.logical_and(r-ii>=0, r-ii<=N) 

output = x[mask].reshape(2,3)

这看起来不错吗？