Замена numpy .apply_along_axis в CuPy - PullRequest
Новая
       64

Замена numpy .apply_along_axis в CuPy

1 голос
/ 08 мая 2020

У меня есть нейронная сеть на основе NumPy, которую я пытаюсь перенести на CuPy. У меня есть следующая функция: 

import numpy as np

def tensor_diag(x): return np.apply_along_axis(np.diag, -1, x)

# Usage: (x is a matrix, i.e. a 2-tensor)
def sigmoid_prime(x): return tensor_diag(sigmoid(x) * (1 - sigmoid(x)))

Это работает с использованием NumPy, но CuPy не имеет аналога для функции (он не поддерживается с 8 мая 2020 года). Как я могу эмулировать это поведение в CuPy?

1 Ответ

1 голос
/ 08 мая 2020 
In [284]: arr = np.arange(24).reshape(2,3,4)

np.diag принимает массив 1d и возвращает 2d со значениями по диагонали. apply_along_axis просто выполняет итерацию по всем измерениям, кроме последнего, и передает последний, по одному массиву, в diag: 

In [285]: np.apply_along_axis(np.diag,-1,arr)                                                          
Out[285]: 
array([[[[ 0,  0,  0,  0],
         [ 0,  1,  0,  0],
         [ 0,  0,  2,  0],
         [ 0,  0,  0,  3]],

        [[ 4,  0,  0,  0],
         [ 0,  5,  0,  0],
         [ 0,  0,  6,  0],
         [ 0,  0,  0,  7]],

        [[ 8,  0,  0,  0],
         [ 0,  9,  0,  0],
         [ 0,  0, 10,  0],
         [ 0,  0,  0, 11]]],


       [[[12,  0,  0,  0],
         [ 0, 13,  0,  0],
         [ 0,  0, 14,  0],
         [ 0,  0,  0, 15]],

        [[16,  0,  0,  0],
         [ 0, 17,  0,  0],
         [ 0,  0, 18,  0],
         [ 0,  0,  0, 19]],

        [[20,  0,  0,  0],
         [ 0, 21,  0,  0],
         [ 0,  0, 22,  0],
         [ 0,  0,  0, 23]]]])
In [286]: _.shape                                                                                      
Out[286]: (2, 3, 4, 4)

Я мог бы сделать то же сопоставление с: 

In [287]: res = np.zeros((2,3,4,4),int)                                                                
In [288]: res[:,:,np.arange(4),np.arange(4)] = arr

проверьте результат apply: 

In [289]: np.allclose(_285, res)                                                                       
Out[289]: True

Или для более прямой копии apply используйте np.ndindex, чтобы сгенерировать все пары кортежей i,j для перебора первых двух габариты arr: 

In [298]: res = np.zeros((2,3,4,4),int)                                                                
In [299]: for ij in np.ndindex(2,3): 
     ...:     res[ij]=np.diag(arr[ij]) 
     ...:                                                                                              
In [300]: np.allclose(_285, res)                                                                       
Out[300]: True
...