在包含多个图像的ndarray中最快的方法是转置和规范化数据？

Question

我有一个“批量”图像，通常是128个，最初被读取为尺寸为128x360x640x3的numpy数组。 我需要将每个图像从NHWC转置为NCHW，因此需要对ndarray.transpose(2,0,1)进行操作， 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只会重新排列步幅，而不会真正触及内存，但是我只为转置（每个120ms）看到大约每张图像约1 ms。 另一方面，同时进行转置和除法操作会使总时间增加到约350ms。 尽可能加快速度的最佳方法是什么？ Cython和多线程处理的组合会有所帮助吗？ 我在Ubuntu上也可以访问OpenMP。

编辑：我尝试了一个简单的multiprocessing.Pool实现，它为整个循环提供了大约270ms，但我想进一步优化它。

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


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

Answer 1

伪数据

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))

如果可以并在整个阵列上进行操作，通常最好避免for循环

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)

Answer 2

您的问题是高度的内存和缓存依赖关系。 最佳解决方案将取决于您的处理器和RAM速度。 这是使用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.5s，这不包括在时序中。