Блог

У меня есть 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 , поэтому я ищу альтернативные решения

Перебор 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]])
  

вдоль 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)
  

Вы действительно должны были привести правильный рабочий пример, сделав его излишне запутанным.

Вот как это сделать в numpy, а теперь подсказка о torch.

Далее выбирается срез длиной n вдоль третьего измерения, начиная с точек idx, в зависимости от двух других измерений:

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

Я придумал ниже, используя широковещательную передачу:

 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)
  

Это выглядит нормально?