[英]Advice on vectorizing block-wise operations in Numpy
我正在尝试实现一系列统计操作,我需要帮助向量化我的代码。
这个想法是从两个图像中提取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 个内核上也需要很长时间.
我的心智模型是 matlab 中的blockproc
是如何工作的 - 我最终以这种模式编写了这段代码。 我将不胜感激有关提高此代码性能的任何建议。
通过使用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)
首先,分析您的代码以确定瓶颈。 您可以使用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
声明:本站的技术帖子网页,遵循CC BY-SA 4.0协议,如果您需要转载,请注明本站网址或者原文地址。任何问题请咨询:yoyou2525@163.com.