使用 NumPy 沿任意维度进行单热编码

One-hot encode along arbitrary dimension with NumPy

给定一个具有任意多个维度的 numpy 数组,我希望能够对这些维度中的任何一个进行一次性编码。例如,假设我有一个形状为 (10, 20, 30, 40) 的数组 a 我可能想对第二个维度进行一次热编码,即转换 a 以便结果仅包含值 0 1a[i, :, j, k]ijk 的每个选择都恰好包含一个零条目(在该维度上前一个最大值的位置).

我考虑过先获取 a.argmax(axis=1) 然后使用 np.ogrid 将其转换为指向最大值的索引,但我无法弄清楚细节。我也担心这种方法的内存消耗。

有没有一种简单的方法可以做到这一点(理想情况下需要很少的额外内存)?

这是 array-assignment -

的一种方式
def onehotencode_along_axis(a, axis):
    # Setup o/p hot encoded bool array 
    h = np.zeros(a.shape,dtype=bool)
    idx = a.argmax(axis=axis)

    # Setup same dimensional indexing array as the input
    idx = np.expand_dims(idx, axis) # Thanks to @Peter

    # Finally assign True values
    np.put_along_axis(h,idx,1,axis=axis)
    return h

样本 运行s 2D 案例 -

In [109]: np.random.seed(0)
     ...: a = np.random.randint(11,99,(4,5))

In [110]: a
Out[110]: 
array([[55, 58, 75, 78, 78],
       [20, 94, 32, 47, 98],
       [81, 23, 69, 76, 50],
       [98, 57, 92, 48, 36]])

In [112]: onehotencode_along_axis(a, axis=0)
Out[112]: 
array([[False, False, False,  True, False],
       [False,  True, False, False,  True],
       [False, False, False, False, False],
       [ True, False,  True, False, False]])

In [113]: onehotencode_along_axis(a, axis=1)
Out[113]: 
array([[False, False, False,  True, False],
       [False, False, False, False,  True],
       [ True, False, False, False, False],
       [ True, False, False, False, False]])

用于验证更高(多维)5D 案例的样本 运行 -

In [114]: np.random.seed(0)
     ...: a = np.random.randint(11,99,(2,3,4,5,6))
     ...: for i in range(a.ndim):
     ...:     out = onehotencode_along_axis(a, axis=i)
     ...:     print np.allclose(out.sum(axis=i),1)
True
True
True
True
True

如果您需要最终输出为具有 0s 和 1s 的 int 数组,请在布尔输出数组上使用视图:

onehotencode_along_axis(a, axis=0).view('i1') 等等。