Какой самый быстрый способ транспонировать и нормализовать данные в массиве, содержащем несколько изображений?

Question

У меня есть «серия» изображений, обычно 128, которые изначально считываются в массив массивов размером 128x360x640x3. Мне нужно переместить каждое изображение из NHWC в NCHW, таким образом, операция ndarray.transpose(2,0,1), а также нормализовать пиксели в диапазоне [0,1], поэтому мне нужно разделить массив на 255. Эта операция пакетной обработки будет повторяться периодически возможно, сто или около того раз. Простейшая реализация для этого выглядит следующим образом:

for i in range(128):
    batchImageDataNew[i,:,:] = batchImageData[i,:,:].transpose(2,0,1)/255.

batchImageDataNew имеет тип np.float32, тогда как batchImageData имеет значение np.uint8. Я стараюсь максимально ускорить этот процесс. Я думал, что ndarray.transpose только переставляет шаги, фактически не затрагивая память, но я вижу приблизительно ~ 1 мс на изображение только для транспонирования (всего 120 мс). С другой стороны, выполнение как транспонирования, так и деления увеличивает общее время до 350 мс. Что было бы лучшим способом максимально ускорить это? Поможет ли сочетание Cython и многопоточности? Я работаю в Ubuntu, где у меня также есть доступ к OpenMP.

РЕДАКТИРОВАТЬ: я попробовал простую многопроцессорную реализацию. Бассейн дал мне около 270 мсек для всего цикла, но я хотел бы оптимизировать его еще дальше.

def preprocess(i):
    batchImageDataNew[i,:,:] = batchImageData[i,:,:].transpose(2,0,1)/255.


pool = multiprocessing.Pool(8)
pool.map(preprocess, range(128))

Answer 1

Ваша проблема сильно зависит от памяти и кэша.Оптимальное решение будет зависеть от вашего процессора и скорости оперативной памяти.Это решение с использованием Numba, но вы можете сделать довольно похожий подход, используя cython.

Пример

import numba as nb
import numpy as np
import time


def tran_scal(batchImageData):
  s=batchImageData.shape
  batchImageDataNew=np.empty((s[0],s[3],s[1],s[2]),dtype=np.float32)
  for i in range(batchImageData.shape[0]):
    batchImageDataNew[i,:,:] = batchImageData[i,:,:].transpose(2,0,1)/255.
  return batchImageDataNew


@nb.njit()
def tran_scal_nb(batchImageData):
  s=batchImageData.shape
  batchImageDataNew=np.empty((s[0],s[3],s[1],s[2]),dtype=np.float32)
  for i in range(batchImageData.shape[0]):
    for j in range(batchImageData.shape[1]):
      for k in range(batchImageData.shape[2]):
        for l in range(batchImageData.shape[3]):
          batchImageDataNew[i,l,j,k] = batchImageData[i,j,k,l]*(1/255.)
  return batchImageDataNew

@nb.njit(parallel=True)
def tran_scal_nb_p(batchImageData):
  s=batchImageData.shape
  batchImageDataNew=np.empty((s[0],s[3],s[1],s[2]),dtype=np.float32)
  for i in nb.prange(batchImageData.shape[0]):
    for j in range(batchImageData.shape[1]):
      for k in range(batchImageData.shape[2]):
        for l in range(batchImageData.shape[3]):
          batchImageDataNew[i,l,j,k] = batchImageData[i,j,k,l]*(1/255.)
  return batchImageDataNew

Время

Core i7-4xxx
#Test data
data=np.array(np.random.rand(128,360,640,3)*255,dtype=np.uint8)
Your solution:    550ms
@wwii(transpose): 379ms
tran_scal_nb:     190ms 
tran_scal_nb_p:   100ms 

При первом вызове накладные расходы на компиляцию составляют около 0,5 с, что не входит в тайминги.

Answer 2

Поддельные данные

a = np.array([[[1,1]],[[2,2]],[[3,3]]])
b = a + 10
c = b + 10
d = c + 10
e = np.stack((a,b,c,d))

Обычно лучше избегать циклов, если вы можете использовать весь массив

f = np.transpose(e, (0,3,1,2))
g = f / 255

>>> e.shape
(4, 3, 1, 2)
>>> f.shape
(4, 2, 3, 1)

или np.moveaxis вместо transpose

f = np.moveaxis(e, 3, 1)
f = np.moveaxis(e, (1,2,3), (2,3,1))

---

Небольшое улучшение ~ 25% (на моей машине) может быть достигнуто путем создания массива заранее, чтобы принять результат деления:

a = np.array(np.random.rand(128,360,640,3)*255,dtype=np.uint8)
b = np.zeros((128,3,360,640), dtype=np.float32)
np.divide(np.moveaxis(a, (1,2,3), (2,3,1)), 255, out=b)