Получить только ненулевые суб-массивы из N-мерного массива Numpy

Question

У меня есть numpy array 'arr', который имеет shape (1756020, 28, 28, 4).В основном 'arr' имеет 1756020 небольших массивов shape (28,28,4).Из 1756020 массивов 967210 имеют значение «все ноль», а 788810 имеет все ненулевые значения.Я хочу удалить все массивы 967210 'all zero'.Я написал цикл if else, используя условие arr[i]==0.any(), но это занимает много времени.Есть ли лучший способ сделать это?

Answer 1

Я сделал небольшой тестовый скрипт с указанным вами размером. На моем компьютере создание массива (ошибка памяти, если значение float, поэтому логические значения) и выбор медленны, но поиск нулей кажется довольно быстрым:

if __name__ == '__main__':
    arr = np.ones((1756020, 28, 28, 4), dtype=bool)
    for i in range(0,1756020,2):
        arr[i] = 0
    print(arr[:5])
    s = arr.shape
    t0 = time.time()
    arr2 = arr.reshape((s[0], np.prod(s[1:])))
    ok = np.any(arr2, axis=1)
    print(time.time()-t0)
    arr_clean = arr2[ok]
    print(time.time()-t0)
    arr_clean = arr_clean.reshape((np.sum(ok), *s[1:]))
    print(time.time()-t0)
    print('end')

Выход:

0.4846000671386719 # Быстрый поиск нулей

29.750200271606445 # Удаление нулей происходит медленно

29.797000408172607 # Преобразовать в исходную форму [1:] быстро

Answer 2

Один из способов векторизации вашей логики - использовать numpy.any с аргументом кортежа для axis, содержащим непроверенные измерения.

# set up 4d array of ones
A = np.ones((5, 3, 3, 4))

# make second of shape (3, 3, 4) = 0
A[1] = 0  # or A[1, ...] = 0; or A[1, :, :, :] = 0

# find out which are non-zero
res = np.any(A, axis=(1, 2, 3))

print(res)

[True False True True True]

Эта функция доступна в numpy v0.17 и выше. Согласно документам :

ось : Нет или int или кортеж целых чисел, необязательно

Если это кортеж целых, сокращение выполняется по нескольким осям, вместо одной оси или всех осей, как раньше.