更换numpy。沿CuPy中的_轴应用_

Question

我有一个基于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没有类似的功能（截至2020年5月8日不支持）。我怎样才能模仿丘比特的这种行为呢

Answer 1

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