python - 在 Numpy Python 中向量化此函数

Question

我有一个包含 60,000 个数字的数组，从 0 到 9:

In [1]: trainY
Out[1]: 
array([[5],
       [0],
       [4],
       ..., 
       [5],
       [6],
       [8]], dtype=int8)

我有一个函数可以将 trainY 中的每个元素转换为 10 元素向量，如下所示:

0 -> [1,0,0,0,0,0,0,0,0,0]
1 -> [0,1,0,0,0,0,0,0,0,0]
2 -> [0,0,1,0,0,0,0,0,0,0]
3 -> [0,0,0,1,0,0,0,0,0,0]
...
9 -> [0,0,0,0,0,0,0,0,0,1]

函数:

def transform_y(y):
    new_y = np.zeros(10)
    new_y[y] = 1
    return new_y

我的代码一次只能处理 1 个元素。一次性转换我的 trainY 数组的最佳方法是什么(for 循环除外)？我应该使用 map 吗？有人还可以告诉我如何重写函数以使其矢量化吗？

谢谢。

Answer 1

您可以大大提高代码速度，创建一个沿对角线排列的二维数组，然后根据输入数组提取正确的行:

a = array([[5],
           [0],
           [4],
           ..., 
           [5],
           [6],
           [8]], dtype=int8)

new_y = np.eye(a.max()+1)[a.ravel()]

一个更快的解决方案是用零创建输出数组，然后根据 a 中的索引填充它:

new_y = np.zeros((a.shape[0], a.max()+1))
new_y[np.indices(a.ravel().shape)[0], a.ravel()] = 1.