Python – извлечение ненулевых пикселей изображения RGBA – ускорение цифровой маски

Анализ

~np.all(f_pts == 0, axis=1) неоптимален по нескольким причинам:

Прежде всего, операция ограничена памятью (обычно пропускной способностью ОЗУ), поскольку изображение огромно (400 МБ в памяти) — даже логические маски огромны (100 МБ). Более того, Numpy создает несколько временных массивов: один для результата f_pts == 0, один для np.all(...) и даже один для ~np.all(...). Каждый массив полностью хранится в оперативной памяти и считывается обратно, что неэффективно. Хуже того: вновь выделенные массивы обычно не отображаются напрямую в физическую память из-за особенностей работы виртуальной памяти, поэтому за каждый временный массив приходится платить дорогостоящие накладные расходы, связанные с ошибками страниц. И последнее, но не менее важное: Numpy не оптимизирован для вычислений массивов, где основное измерение (т. е. обычно последнее) крошечное, как здесь (4 элемента). Однако это можно исправить, используя лучшую структуру памяти. См. AoS против SoA для получения дополнительной информации.

Обратите внимание, что np.zeros резервирует некоторое обнуленное пространство в виртуальной памяти, но не отображает его физически. Таким образом, первый вызов f_img_to_pts выполняется медленнее из-за ошибок страницы. В реальных приложениях это, как правило, не так, поскольку запись нулей в img или однократное чтение приводит к физическому отображению всей страницы (т. е. к ошибкам страницы). Если это не так, то вам обязательно следует использовать разреженные матрицы вместо плотных. Предположим, img уже отображен в физической памяти (это можно сделать, запустив код дважды или просто вызвав img.fill(0) после np.zeros.

Также обратите внимание, что dtype float на самом деле является 64-битным типом с плавающей запятой, а не 32-битным. 64-битные числа с плавающей запятой обычно дороже, особенно в кодах, привязанных к памяти, поскольку они занимают вдвое больше места. Вам обязательно следует использовать 32-битные для вычислений изображений, поскольку они почти никогда не требуют такой очень высокой точности. Предположим, что входной массив теперь 32-битный.

Решения

Один из способов уменьшить накладные расходы — это итерация строка за строкой (циклы не так уж и плохи, как здесь кажется, а на самом деле совсем наоборот, если все сделано правильно). Хотя это должно быть быстрее, это все равно неоптимально. Невозможно сделать это намного быстрее из-за особенностей конструкции Numpy.

Чтобы сделать это еще быстрее, можно использовать модули, компилирующие функции Python в нативные функции, такие как Numba или Cython. С их помощью мы можем легко избежать создания временных массивов, специализировать код для 4-канальных изображений и даже использовать несколько потоков. Вот полученный код:

import numpy as np
import time
import numba as nb

img_size = (10000, 10000)
img = np.zeros((*img_size, 4), np.float32)  # Make RGBA image
img.fill(0)

# Put some values for pixels [float, float, float, int]
img[0][1] = [1.1, 2.2, 3.3, 4]
img[1][0] = [0, 0, 0, 10]
img[1][2] = [6.1, 7.1, 8.1, 0]

# Compute ~np.all(f_pts==0,axis=1) and assume most items are 0
@nb.njit('(float32[:,:,::1],)', parallel=True)
def f_img_to_pts_optim(f_img):  # Get non-zero rows with values from whole img array
    n, m, c = f_img.shape
    assert c == 4 # For sake of performance
    f_mask = np.zeros(n * m, dtype=np.bool_)
    f_pts = np.reshape(f_img, (n*m, c))
    for ij in nb.prange(n * m):
        f_mask[ij] = (f_pts[ij, 0] != 0) | (f_pts[ij, 1] != 0) | (f_pts[ij, 2] != 0) | (f_pts[ij, 3] != 0)
    f_pts = f_pts[f_mask]
    return f_pts

t1 = time.time()
pxs = f_img_to_pts_optim(img)
t2 = time.time()
print('PIXELS EXTRACTING TIME: ', t2 - t1)
print(pxs)

Полученные результаты

Вот тайминги на 10-ядерном процессоре Skylake Xeon:

Initial code (64-bit):          1720 ms
Initial code (32-bit):          1630 ms
Optimized code (sequential):     323 ms
Optimized code (parallel):       182 ms   <----------

Оптимизированный код в 9 раз быстрее исходного кода (при том же типе ввода). Обратите внимание, что код не масштабируется в зависимости от количества ядер, поскольку чтение входных данных является узким местом в сочетании с последовательной операцией f_pts[f_mask]. Можно построить f_pts[f_mask], используя несколько потоков (и это должно быть примерно в 2 раза быстрее), но это довольно сложно сделать (особенно с Numba).

Обычно я «прибегаю» к C++, когда мне нужно ускорить код Python. Однако ответ Жерома Ришара меня впечатлил, особенно использование numba, с которым я не был знаком. В качестве альтернативы, вот параллельное решение с использованием OpenMP, которое обеспечивает пятикратное ускорение по сравнению с подходом на основе JIT (на i7 11700K). Он использует тип __int128, доступный в GCC, для выполнения нулевых тестов для четырех чисел с плавающей запятой одновременно.

Сторона Python:

    #!/usr/bin/env python3
    
    import numpy as np
    from ctypes import *
    from time import time
    import numba as nb
    
    from numpy.ctypeslib import ndpointer
    from math import prod
    
    clib = CDLL('./libfilter.so')
    clib.FindNZ_IndexesF32x4.restype  = c_uint
    clib.FindNZ_IndexesF32x4.argtypes = (c_void_p, # inp: int128 (see f_img_to_pts_cpp(..))
                                         ndpointer(c_uint,
                                                   flags = "C_CONTIGUOUS"), # indexes
                                         c_uint) # size
    
    
    def f_img_to_pts_cpp(inp):
        cpp   = 4 # components per pixel
        assert inp.shape[-1] == cpp
    
        flat_size = prod(inp.shape)
        inp       = inp.reshape(flat_size)
    
        # The ctypes module doesn't have a 128-bit type, so we manually assert
        # the preconditions for safe casting
        # (except alignment, but it works as it is, at least on x86_64)
        assert inp.flags['C_CONTIGUOUS']
        assert inp.dtype == np.float32
    
        inp     = inp.view(np.complex128) # cast to 128-bit elements
        indexes = np.empty(flat_size, dtype=np.uint32) # uninitialzed array
        size    = clib.FindNZ_IndexesF32x4(inp.ctypes.data, 
                                           indexes,
                                           flat_size//cpp)
        indexes = indexes[:size]
        out = inp[indexes].view(np.float32) # cast back to float
        return out.reshape((size, cpp))
    
    
    @nb.njit('(float32[:,:,::1],)', parallel=True)
    def f_img_to_pts_optim(f_img):  # Get non-zero rows with values from whole img array
        n, m, c = f_img.shape
        assert c == 4 # For sake of performance
        f_mask = np.zeros(n*m, dtype=np.bool_)
        f_pts = np.reshape(f_img, (n*m, c))
        for ij in nb.prange(n*m):
            f_mask[ij] = (f_pts[ij, 0] != 0) | (f_pts[ij, 1] != 0) | \
                (f_pts[ij, 2] != 0) | (f_pts[ij, 3] != 0)
        f_pts = f_pts[f_mask]
        return f_pts
    
    
    def py_opt(img):
        print('warming up...') # "spin up" the processor 
        for i in range(4):
            f_img_to_pts_optim(img)
    
        tstart = time()
        ret = f_img_to_pts_optim(img)
        runtime = time() - tstart
        print(f'numba runtime {runtime:.4f}')
        return ret
    
    def cpp_opt(img):
        tstart = time()
        pxs = f_img_to_pts_cpp(img)
        runtime = time() - tstart
        print(f'C++ runtime {runtime:.4f}')
        return pxs
    
    img = np.zeros((10000, 10000, 4), np.float32)
    
    img[0][1] = [1.1, 2.2, 3.3, 4]
    img[1][0] = [0,   0,   0,   10]
    img[1][2] = [6.1, 7.1, 8.1, 0]
    
    py_pxs  = py_opt(img)
    cpp_pxs = cpp_opt(img)
    
    print(py_pxs)
    assert (py_pxs == cpp_pxs).all()

Сторона С++:

#include <cstdint>
#include <atomic>

using u32    = uint32_t;
using int128 = __int128;

extern "C"
u32 FindNZ_IndexesF32x4(const int128 *  inp,
                        u32          *  indexes,
                        const u32 size) noexcept {
  std::atomic_uint l(0);

  #pragma omp parallel for
  for(u32 i = 0; i < size; i++) {
    if (inp[i]) {
      u32 j = l.fetch_add(1);
      indexes[j] = i;
    }
  }

  return l;
}

Чтобы построить исходный код выше:

g++ -DNDEBUG -fPIC -Ofast -mtune=native -fopenmp -std=gnu++17 -shared -o libfilter.so filter.cpp

Выход:

warming up...
numba runtime 0.1169
C++ runtime 0.0163
[[ 1.1  2.2  3.3  4. ]
 [ 0.   0.   0.  10. ]
 [ 6.1  7.1  8.1  0. ]]

РЕДАКТИРОВАТЬ

Как правильно заметил @Jérôme Richard, недетерминированная природа потоков означает, что порядок индексов (и, следовательно, порядок результата) может отличаться от варианта JIT. Однако, поскольку количество ненулевых элементов крайне мало, сортировка индексов практически не требует затрат.

Согласование: согласно этому Почему malloc выровнен по 16 байтам?, Адреса, выделенные с помощью malloc, в системе x86_64 выравниваются по 16 байтам, если только запрошенный размер не меньше 15 байт, что не имеет место в данном контексте.

#include <cstdint>
#include <atomic>
#include <algorithm>

using u32    = uint32_t;
using int128 = __int128;

extern "C"
u32 FindNZ_IndexesF32x4(const int128 *  inp,
                        u32          *  indexes,
                        const u32 size) noexcept {
  std::atomic_uint l(0);

  // It's always aligned on a x86_64 system
  // const bool is_misaligned  = (size_t(inp) & 15) != 0;

  // Regarding -0.0:
  int128    mask;
  {
    u32 * u = (u32*)&mask;
    u[0] = (1u << 31) - 1;
    u[1] = u[0];
    u[2] = u[0];
    u[3] = u[0];
  }

  #pragma omp parallel for
  for(u32 i = 0;  i < size;  i++) {
    if (inp[i] & mask) {
      u32 j = l.fetch_add(1);
      indexes[j] = i;
    }
  }

  std::sort(indexes, indexes + l);

  return l;
}

Вывод (с маскированием и сортировкой, i7 11700K):

warming up...
numba runtime 0.142
C++ runtime 0.0171
[[ 1.1  2.2  3.3  4. ]
 [ 0.   0.   0.  10. ]
 [ 6.1  7.1  8.1  0. ]]

Но я бы согласился, что мое решение несколько переоснащено. Функция JIT более универсальна и будет хорошо работать даже со значительным количеством ненулевых элементов.

ОБНОВЛЯТЬ

@Jérôme Richard сделал еще одно наблюдение: эффективная пропускная способность памяти в моих результатах кажется слишком хорошей, чтобы быть правдой. Мне потребовалось некоторое тестирование и поиск в Google, чтобы понять, что происходит. Рассмотрим эту функцию:

...
extern "C"
u64 CPP_Reduce(const u64 * inp, const u32 size) noexcept {
  u64 acc = 0;
  #pragma omp parallel for reduction(+: acc)
  for(u32 i = 0;  i < size;  i++)
    acc += inp[i];
  return acc;
}
...

Давайте передадим img из Python:

    clib.CPP_Reduce.argtypes = (ndpointer(c_uint64,
                                          flags = "C_CONTIGUOUS"),
                                c_uint) # size
    clib.CPP_Reduce.restype = c_uint64
...
    as_vec64 = img.reshape(prod(img.shape)).view(dtype=np.uint64)
    print(f'reduction input: size {as_vec64.nbytes}  shape {as_vec64.shape}')

    n_iters = 10
    tstart = time()
    for i in range(n_iters):
        clib.CPP_Reduce(as_vec64, as_vec64.shape[0])
    dt = (time() - tstart) / n_iters
    print(f'dt: {dt:.4f} s')
    bandwidth = as_vec64.nbytes / dt / 2**30
    print(f'bandwidth: {bandwidth:.4f} GB/s')
...

Выход:

reduction input: size 1600000000  shape (200000000,)
dt: 0.0123 s
bandwidth: 120.8112 GB/s

Это i5 1335U с тактовой частотой оперативной памяти 4,267 ГГц. Общая ширина шины составляет 128 бит. Таким образом, его пиковая пропускная способность составляет 4,267e9 * 16/2**30 = 63,58 ГБ/с. На первый взгляд может показаться, что пример приведения некорректен, но это не так (надеюсь). Давайте запустим нашу функцию сокращения с другим входом:

img = np.arange(0, 10000*10000*4, 1, np.float32).reshape((10000, 10000, 4))

Выход:

reduction input: size 1600000000  shape (200000000,)
dt: 0.0458 s
bandwidth: 32.5458 GB/s

Этот результат гораздо менее удивителен (и, вероятно, не очень хорош, но это не по теме). Вы можете прочитать объяснение здесь: Ядро Linux: роль нулевого выделения страниц во время paging_init По сути, из-за трюков с подкачкой img может эффективно полностью поместиться в кэше ЦП, если он в основном равен нулям, но в противном случае «ведет себя так, как ожидалось».