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

along the 3D axis:

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

by enumerating you get the index of the axis and index from where you want to start slicing in one.

narrow gives you a zero-copy length long slice from a starting index start along a certain axis

you said you wanted:

dim = 2
start = index
length = 2

then you simply have to stack these tensors back to a single 3D.

This is the least work intensive thing i can think of for pytorch.

EDIT

if you just want different indices along different axis and indices is a 2D tensor you can do:

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

You really should have given a proper working example, making it unnecessarily confusing.

Answer 3

Here is how to do it in numpy, now clue about torch, though.

The following picks a slice of length n along the third dimension starting from points idx depending on the other two dimensions:

# 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)

这看起来还好吗？

跨特定维度的明确切片

4 个答案: