关于在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秒才能完成，我并不过分担心预处理时间。下面的步骤是我想要优化的步骤

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

@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 x 10Kpx图像发送到cauchy_schwartz函数，它将在不到一秒钟的时间内处理所有内容，但使用我的方法，即使在4核上，也需要很长时间

我的心智模型是matlab中的blockproc是如何工作的——我最终以这种模式编写了这段代码。如果您能给我一些关于改进此代码性能的建议，我将不胜感激

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。由于您不太关心预处理时间，因此假设您已获得包含展平面片的数组params

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

正如您可以看到的params的形状是(3,2,100)，三个数字3、2和100是随机选择的，以创建一个辅助数组来演示使用apply_along_axis的逻辑。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)