np.where在多维数组的应用

函数用途

返回查找的参数，在数组中的索引。

Code

举例：一般卷积神经网络的输入或者输出为一个四维的数组/Tensor。一般为[batch_size, channel, height, width]，下面代码目标是输出所有值为0的数字的索引。

output = [[
    [[1, 0, 2],
    [2, 1, 0],
    [1, 0, 0]]
]]

arr = np.array(output)
print(arr.shape)
res = np.where(arr==0)
print(res)

Output

# print(arr.shape)
(1, 1, 3, 3)
# print(res)
(array([0, 0, 0, 0], dtype=int64), array([0, 0, 0, 0], dtype=int64), array([0, 1, 2, 2], dtype=int64), array([1, 2, 1, 2], dtype=int64))

np.where的输出结果为一个list，里面包含4个ndarray，分别代表四维。

[0, 0, 0, 0] # axis=0
[0, 0, 0, 0] # axis=1
[0, 1, 2, 2] # axis=2
[1, 2, 1, 2] # axis=3

正确读值，从列来看，四个0值的索引分别是

print(arr[0][0][0][1]) # output:0
print(arr[0][0][1][2]) # output:0
print(arr[0][0][2][1]) # output:0
print(arr[0][0][2][2]) # output:0