如何更快/并行化将多个Numpy 3D数组合并至已有3D数组？

问题背景 #

我有一个包含N个形状为(H, W, 4)的NumPy数组列表（示例中N=100，H=W=1024），同时预创建了一个形状为(H, W*N, 4)的NumPy数组imgs_arr（仅启动时创建一次，后续复用）。当前通过循环逐个将列表中的数组填充到imgs_arr中，代码如下：

import numpy as np

imgs = [np.random.randint(0, 255, (1024, 1024, 4), np.uint8) for _ in range(100)]
imgs_arr = np.empty((1024, 1024 * 100, 4), np.uint8)
for i in range(100):
    imgs_arr[:, i * 1024:(i + 1) * 1024] = imgs[i]

该操作耗时0.1025秒。

已尝试方案及问题 #

我曾尝试用Numba并行优化，但效果有限且引入额外开销：

import numpy as np

from numba import carray, njit, prange
from numba.extending import intrinsic
from numba.typed import List

@intrinsic
def address_to_void_pointer(typingctx, src):
    from numba.core import types, cgutils
    sig = types.voidptr(src)

    def codegen(cgctx, builder, sig, args):
        return builder.inttoptr(args[0], cgutils.voidptr_t)

    return sig, codegen


@njit(cache=True, parallel=True)
def numba_parallelization_func(imgs_arr_addr, imgs, shape,
                               dtype, img_count):
    imgs_arr = carray(address_to_void_pointer(imgs_arr_addr), shape, dtype)
    for i in prange(img_count):
        imgs_arr[:, i * 1024:(i + 1) * 1024] = imgs.getitem_unchecked(i)

    return None


def numba_parallelization() -> None:
    imgs = List([np.random.randint(0, 255, (1024, 1024, 4), np.uint8) for _ in range(100)])
    imgs_arr = np.empty((1024, 1024 * 100, 4), np.uint8)
    imgs_arr_addr = imgs_arr.ctypes.data
    numba_parallelization_func(imgs_arr_addr, imgs, imgs_arr.shape,
                               imgs_arr.dtype, len(imgs))

    return None

编译/缓存后numba_parallelization_func耗时0.0975秒，仅略快于原循环；且需要将普通列表转为Numba Typed List，普通列表append循环耗时5.4999e-06秒，但预分配Typed List并赋值的循环耗时0.2541秒，直接抵消了Numba的性能提升。

补充说明 #

实际场景中，数组数量可达数万，且图片通过缓存生成，多数请求中仅需加载部分图片，无法将所有图片存入内存。

优化方案 #

1. NumPy原生concatenate（一次性全量填充） #

如果是一次性将所有列表中的数组合并到目标数组，直接使用np.concatenate，利用NumPy内部的内存块复制优化，性能远高于Python循环：

imgs_arr[:, :] = np.concatenate(imgs, axis=1)

该操作跳过Python循环的解释器开销，直接在底层完成连续内存复制，耗时会显著低于0.1秒。

2. 分块线程池填充（部分更新场景） #

如果每次仅需更新imgs_arr的部分区域，可将任务分块，用线程池处理（内存复制属于带宽受限操作，线程池开销远低于多进程）：

from concurrent.futures import ThreadPoolExecutor
import numpy as np

# 启动时初始化线程池（仅执行一次）
executor = ThreadPoolExecutor(max_workers=4)

def fill_chunk(args):
    img_idx, img, target_arr = args
    chunk_width = img.shape[1]
    start_col = img_idx * chunk_width
    end_col = start_col + chunk_width
    target_arr[:, start_col:end_col] = img

# 填充时提交任务
tasks = [(i, imgs[i], imgs_arr) for i in range(len(imgs))]
executor.map(fill_chunk, tasks)

注意：线程数建议设为CPU核心数，过多线程会引发调度开销，反而降低性能。

3. 简化版Numba实现（避免Typed List） #

若坚持使用Numba，可直接传递普通列表，无需转换为Typed List，消除额外开销：

import numpy as np
from numba import njit, prange

@njit(parallel=True, cache=True)
def numba_fill(target_arr, imgs_list):
    num_imgs = len(imgs_list)
    chunk_width = target_arr.shape[1] // num_imgs
    for i in prange(num_imgs):
        target_arr[:, i*chunk_width:(i+1)*chunk_width] = imgs_list[i]

# 直接传入普通列表调用
numba_fill(imgs_arr, imgs)

该版本编译后性能略优于原Python循环，且没有列表转换的额外耗时。

4. 内存布局优化 #

确保目标数组和子数组均为连续内存布局，可大幅提升内存复制效率：

# 预处理子数组为连续内存
imgs = [np.ascontiguousarray(img) for img in imgs]
# 确保目标数组为连续内存
imgs_arr = np.ascontiguousarray(imgs_arr)

内容的提问来源于stack exchange，提问作者user-cd