Python 在包含多个图像的数据阵列中，最快的转置和规范化数据的方法是什么？

我有一批图像，通常是128张，最初被读取到一个128x360x640x3的numpy数组中。我需要将每个图像从NHWC转置到NCHW，从而执行

ndarray.transpose（2,0,1）

操作，并将像素规格化为[0,1]范围，因此我需要将数组除以255。该批处理操作将定期重复大约100次。最简单的实现如下所示：

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

编辑：我尝试了一个简单的多处理池实现，它为整个循环提供了大约270ms的时间，但我想进一步优化它

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


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

假数据

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)

假数据

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)

---

您的问题高度依赖于内存和缓存。最佳解决方案取决于处理器和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.5秒，这不包括在计时中。

您的问题高度依赖于内存和缓存。最佳解决方案取决于处理器和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，这不包括在计时中。

我正在使用Jetson TX2（ARM处理器），因此它并不是最大的。我将测试您的解决方案：Cython会比这更快吗（例如更接近纯C）？这都是关于缓存的使用。我不知道，在你的平台上哪一个是最快的，但我会在Cython中尝试同样的。“Cython或更接近纯C示例”是什么意思？对不起，我的意思是：我想知道Cython的性能是否会接近人们在纯C中所期望的性能。我一定会尝试一下。我很好奇：看我的编辑，你的计时是否也有同样的改进？@wwii我看到性能提升了11%。如果我将内存分配放在基准测试循环之外（仅用于测试，只有当您有内存缓冲区并逐个处理图像时，这才有意义），那么您的函数将有45%的改进。我测试我的tran_scal_nb_p时也在性能测试循环之外分配内存（给出50%）。如果不测试内存分配的话，40-50%的内存分配不会给我带来太多的支持（内存分配非常昂贵），但是在您编辑的示例中，11-25%的内存分配确实可以。也许是垃圾收集器的怪异行为？我正在使用Jetson TX2（ARM处理器），所以它并不是最棒的。我将测试您的解决方案：Cython会比这更快吗（例如更接近纯C）？这都是关于缓存的使用。我不知道，在你的平台上哪一个是最快的，但我会在Cython中尝试同样的。“Cython或更接近纯C示例”是什么意思？对不起，我的意思是：我想知道Cython的性能是否会接近人们在纯C中所期望的性能。我一定会尝试一下。我很好奇：看我的编辑，你的计时是否也有同样的改进？@wwii我看到性能提升了11%。如果我将内存分配放在基准测试循环之外（仅用于测试，只有当您有内存缓冲区并逐个处理图像时，这才有意义），那么您的函数将有45%的改进。我测试我的tran_scal_nb_p时也在性能测试循环之外分配内存（给出50%）。如果不测试内存分配的话，40-50%的内存分配不会给我带来太多的支持（内存分配非常昂贵），但是在您编辑的示例中，11-25%的内存分配确实可以。也许是垃圾收集器的怪异行为？