关于在 Numpy 中向量化块操作的建议

Question

我正在尝试实施一系列统计操作，我需要矢量化我的代码的帮助。

想法是从两个图像中提取 NxN 个补丁，计算这两个补丁之间的距离度量。

为此，我首先使用以下循环构造补丁：

params = []
for i in range(0,patch1.shape[0],1):
    for j in range(0,patch1.shape[1],1):
        window1 = np.copy(imga[i:i+N,j:j+N]).flatten()
        window2 = np.copy(imgb[i:i+N,j:j+N]).flatten()
        params.append((window1, window2))
print(f"We took {time()- t0:2.2f} seconds to prepare {len(params)/1e6} million patches.")

这大约需要 10 秒才能完成，我并不太关心预处理时间。接下来的步骤是我要优化的步骤。

在此之后，为了加快处理速度，我使用了 multipool 来计算实际结果。包含实际计算的函数如下：

@njit
def cauchy_schwartz(imga, imgb):
    p, _ = np.histogram(imga, bins=10)
    p = p/np.sum(p)
    q, _ = np.histogram(imgb, bins=10)
    q = q/np.sum(q)

    n_d = np.array(np.sum(p * q)) 
    d_d = np.array(np.sum(np.power(p, 2) * np.power(q, 2)))
    return -1.0 * np.log10( n_d, d_d)

我用这个结构来处理所有的补丁：

def f(param):
    return cauchy_schwartz(*param)

with Pool(4) as p:
    r = list(tqdm.tqdm(p.imap(f,params), total=len(params)))

我相信一定有更优雅的方法可以做到这一点，因为如果我将整个 10Kpx x 10Kpx 图像发送到 cauchy_schwartz 函数，它会在一秒钟内处理所有内容，但使用我的方法，即使在 4 核上也需要很长时间。

我的心智模型是 blockproc 在 matlab 中的工作方式 - 我最终以该模式编写了这段代码。如果有任何关于提高此代码性能的建议，我将不胜感激。

Answer 1

首先，分析您的代码以确定瓶颈。您可以使用 https://mg.pov.lt/profilehooks/。我认为瓶颈在于补丁的创建，因为您正在为流程创建补丁的副本。您可以通过仅传递补丁的索引来使用更少的内存：

params = []
for i in range(0,patch1.shape[0],1):
    for j in range(0,patch1.shape[1],1):
        start, end = (i,i+N), (j,j+N)
        params.append((start, end))

然后，假设 imga 和 imgb 是全局的，您可以从 cauchy_schwartz 函数创建补丁，如下所示：

@njit
def cauchy_schwartz(start, end):

    a,b = start; c,d = end
    window1 = np.copy(imga[a:b, c:d]).flatten()
    window2 = np.copy(imgb[a:b, c:d]).flatten()

    # process patches window1 and window2

Answer 2

通过使用 apply_along_axis，您可以摆脱 cauchy_schwartz。由于您并不过分关心 pre-processing 时间，因此假设您已获得包含扁平化补丁

的数组 params

params = np.random.rand(3,2,100)

可以看到params的形状是(3,2,100)，随机选择3、2、100这三个数字来创建一个辅助数组来演示使用[=的逻辑12=]。 3 对应于您拥有的补丁数（由补丁形状和图像大小决定），2 对应于两个图像，100 对应于展平的补丁。因此，params 的轴是 (idx of patches, idx of images, idx of entries of a flattened patch)，这与您的代码

创建的列表 params 完全匹配

params = []
for i in range(0,patch1.shape[0],1):
    for j in range(0,patch1.shape[1],1):
        window1 = np.copy(imga[i:i+N,j:j+N]).flatten()
        window2 = np.copy(imgb[i:i+N,j:j+N]).flatten()
        params.append((window1, window2))

使用辅助数组params，这是我的解决方案：

hist = np.apply_along_axis(lambda x: np.histogram(x,bins=11)[0],2,params)
hist = hist / np.sum(hist,axis=2)[...,None]

n_d = np.sum(np.product(hist,axis=1),axis=1)
d_d = np.sum(np.product(np.power(hist,2),axis=1),axis=1)
res = -1.0 * np.log10(n_d, d_d)

关于在 Numpy 中向量化块操作的建议

Advice on vectorizing block-wise operations in Numpy

python

performance

numpy

scipy

scientific-computing