在主环路外广播加速矢量化的小天体行动？

Question

我正在用numpy做一些向量代数，我的算法的挂钟性能看起来很奇怪。程序大致如下：

1. 创建三个矩阵：Y（KxD），X（NxD），T（KxN）
2. 对于Y的每一行：
3. 从X的每一行中减去Y[i]（通过广播）
4. 沿一个轴求差的平方，求和，取平方根，然后存储在T。你知道吗

然而，根据我如何执行广播，计算速度有很大的不同。考虑代码：

import numpy as np
from time import perf_counter

D = 128
N = 3000
K = 500

X = np.random.rand(N, D)
Y = np.random.rand(K, D)
T = np.zeros((K, N))

if True: # negate to enable the second loop
    time = 0.0
    for i in range(100):
        start = perf_counter()
        for i in range(K):
            T[i] = np.sqrt(np.sum(
                np.square(
                  X - Y[i] # this has dimensions NxD
                ),
                axis=1
            ))
        time += perf_counter() - start
    print("Broadcast in line: {:.3f} s".format(time / 100))
    exit()

if True:
    time = 0.0
    for i in range(100):
        start = perf_counter()
        for i in range(K):
            diff = X - Y[i]
            T[i] = np.sqrt(np.sum(
                np.square(
                  diff
                ),
                axis=1
            ))
        time += perf_counter() - start
    print("Broadcast out:     {:.3f} s".format(time / 100))
    exit()

每个循环的时间是单独测量的，平均执行100次。结果是：

Broadcast in line: 1.504 s
Broadcast out:     0.438 s

唯一的区别是第一个循环中的广播和减法是在线完成的，而在第二种方法中，我是在任何向量化操作之前完成的。为什么会有这么大的不同？你知道吗

我的系统配置：

• 联想ThinkStation P920，2x至强银4110，64 GB内存
• 徐邦图18.04.2 LTS（仿生）
• Python 3.7.3（GCC 7.3.0）
• NUMPY1.16.3与OpenBLAS链接（这是np.__config__.show()告诉我的）

---

PS：是的，我知道这可能会进一步优化，但现在我想了解引擎盖下发生了什么。你知道吗

Answer 1

这不是广播问题

我还添加了一个优化的解决方案，以查看实际计算需要多长时间，而不需要大量的内存分配和释放开销。你知道吗

功能

import numpy as np
import numba as nb

def func_1(X,Y,T):
    for i in range(K):
        T[i] = np.sqrt(np.sum(np.square(X - Y[i]),axis=1))
    return T

def func_2(X,Y,T):
    for i in range(K):
        diff = X - Y[i]
        T[i] = np.sqrt(np.sum(np.square(diff),axis=1))
    return T

@nb.njit(fastmath=True,parallel=True)
def func_3(X,Y,T):
    for i in nb.prange(Y.shape[0]):
        for j in range(X.shape[0]):
            diff_sq_sum=0.
            for k in range(X.shape[1]):
                diff_sq_sum+= (X[j,k] - Y[i,k])**2
            T[i,j]=np.sqrt(diff_sq_sum)
    return T

计时

我把所有的计时都放在一个笔记本上，观察到一个非常奇怪的行为。以下代码位于一个单元格中。我还多次尝试调用timit，但在第一次执行单元格时，这并没有改变任何事情。你知道吗

单元格的第一次执行

D = 128
N = 3000
K = 500

X = np.random.rand(N, D)
Y = np.random.rand(K, D)
T = np.zeros((K, N))

#You can do it more often it would not change anything
%timeit func_1(X,Y,T)
%timeit func_1(X,Y,T)

#You can do it more often it would not change anything
%timeit func_2(X,Y,T)
%timeit func_2(X,Y,T)

###Avoid measuring compilation overhead###
%timeit func_3(X,Y,T)
##########################################
%timeit func_3(X,Y,T)

774 ms ± 6.81 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
768 ms ± 2.88 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
494 ms ± 2.09 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
494 ms ± 1.06 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
10.7 ms ± 1.25 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
6.74 ms ± 39.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

第二次执行

345 ms ± 16.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
337 ms ± 3.72 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
322 ms ± 834 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)
323 ms ± 1.15 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
6.93 ms ± 234 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
6.9 ms ± 87.4 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)