Python：如何加速这个功能并使其更具可扩展性？

Question

我有以下函数，它接受形状为 (20,000 x 20,000) 的指示矩阵。我必须运行该函数 20,000 x 20,000 = 400,000,000 次。请注意，indicator_Matrix当作为参数传递给函数时，必须采用 pandas 数据帧的形式，因为我实际问题的数据帧具有 timeIndex 和整数列，但为了理解问题，我对此进行了一些简化。

熊猫实施

indicator_Matrix = pd.DataFrame(np.random.randint(0,2,[20000,20000]))
def operations(indicator_Matrix):
   s = indicator_Matrix.sum(axis=1)
   d = indicator_Matrix.div(s,axis=0)
   res = d[d>0].mean(axis=0)
   return res.iloc[-1]

我尝试通过使用来改进它numpy，但它仍然需要很长时间才能运行。我也尝试过concurrent.future.ThreadPoolExecutor，但仍然需要很长时间才能运行，并且列表理解没有太大改进。

Numpy 实现

indicator_Matrix = pd.DataFrame(np.random.randint(0,2,[20000,20000]))
def operations(indicator_Matrix):
   s = indicator_Matrix.to_numpy().sum(axis=1)
   d = (indicator_Matrix.to_numpy().T / s).T
   d = pd.DataFrame(d, index = indicator_Matrix.index, columns = indicator_Matrix.columns)
   res = d[d>0].mean(axis=0)
   return res.iloc[-1]

output = [operations(indicator_Matrix) for i in range(0,20000**2)]

请注意，我再次转换为数据框的原因d是因为我需要获取列均值并仅保留最后一列均值，使用.iloc[-1]. d[d>0].mean(axis=0)返回列的意思，即

2478    1.0
0       1.0

更新：我仍然陷入这个问题。cudf我想知道在我的本地桌面上使用像和那样的 GPU 软件包CuPy是否会产生任何影响。

Answer 1

你正在做一些不必要的额外数学计算。用简单的英语来说，你正在做的是：

\n

\n
1. 对每列求和
\n
2. 将总和列表“横向”转动并将每列除以它
\n
3. 取每列的平均值，忽略值 \xe2\x89\xa4 0
\n
4. 仅返回最右边的均值
\n

\n

第一步之后，除了最右边的列之外，您不再需要任何东西；您可以忽略其他列，只对您关心的结果进行除法和平均。相应地更改您的代码：

\n

def operations_simpler(indicator_matrix):\n    sums = indicator_matrix.sum(axis=1)\n    last_column = indicator_matrix.iloc[:, -1]\n    divided = last_column / sums\n    return divided[divided > 0].mean()\n

\n

...产生相同的结果，并且花费大约百分之一的时间。从较短的测试运行推断，这将我的机器上 400,000,000 次运行的时间从大约 114 年缩短到......大约 324 天。还是不太好。到目前为止，我还没有通过转换为 NumPy、使用 Numba 编译或使用多重处理来让它运行得更快，但我会继续发布它，以防它有帮助。

\n

注意：您不太可能从线程中看到像这样的计算密集型工作的任何改进；如果有的话，你会想使用多处理。concurrent.futures为两者提供执行者。线程对于避免等待 I/O 非常有用。

\n

Answer 2

假设 @CrazyChucky 的答案是正确的，则可以实现更快的并行 Numba 实现。这个想法是使用普通循环并关心以连续方式读取数据。连续读取数据很重要，因此可以使计算缓存友好/内存高效。这是一个实现：

import numba as nb

@nb.njit(['(int_[:,:],)', '(int_[:,::1],)', '(int_[::1,:],)'], parallel=True)
def compute_fastest(matrix):
    n, m = matrix.shape
    sum_by_row = np.zeros(n, matrix.dtype)
    is_row_major = matrix.strides[0] >= matrix.strides[1]
    if is_row_major:
        for i in nb.prange(n):
            s = 0
            for j in range(m):
                s += matrix[i, j]
            sum_by_row[i] = s
    else:
        for chunk_id in nb.prange(0, (n+63)//64):
            start = chunk_id * 64
            end = min(start+64, n)
            for j in range(m):
                for i2 in range(start, end):
                    sum_by_row[i2] += matrix[i2, j]
    count = 0
    s = 0.0
    for i in range(n):
        value = matrix[i, -1] / sum_by_row[i]
        if value > 0:
            s += value
            count += 1
    return s / count

# output = [compute_fastest(indicator_Matrix.to_numpy()) for i in range(0,20000**2)]

Pandas 数据帧可以包含行优先和列优先数组。关于内存布局，最好对行或列进行迭代。这就是为什么有两种基于 的 sum 实现的原因is_row_major。还有 3 个 Numba 签名：一种用于行主连续数组，一种用于列主连续数组，一种用于非连续数组。Numba 将编译 3 个函数变体并在运行时自动选择最佳的一个。当已知输入 2D 数组是连续的时，Numba 的 JIT 编译器可以生成更快的实现（例如使用 SIMD 指令）。

实验结果

此计算速度比我的 i5-9600KF 处理器（6 核）快约 14.5 倍operations_simpler。它仍然需要很多时间，但计算是受内存限制的，并且在我的机器上几乎是最佳的：它受必须读取的主内存的限制：

On a 2000x2000 dataframe with 32-bit integers:
 - operations:           86.310       ms/iter
 - operations_simpler:    5.450       ms/iter
 - compute_fastest:       0.375       ms/iter
 - optimal:               0.345-0.370 ms/iter

如果你想获得更快的代码，那么你需要使用更紧凑的数据类型。例如，uint8数据类型足够大，可以包含值 0 和 1，但在 Windows 上，它在内存中的大小要小 4 倍。这意味着在这种情况下代码速度可以提高 4 倍。数据类型越小，程序速度越快。人们甚至可以尝试使用位调整将 8 列压缩为 1 列，尽管使用 Numba 通常会慢得多，除非您有大量可用核心。

笔记与讨论

上面的代码仅适用于统一类型的列。如果不是这种情况，您可以将数据帧拆分为多个组，并将每个列组转换为 Numpy 数组，以便调用 Numba 函数（修改为支持组）。请注意，@CrazyChucky 代码也有类似的问题：将混合数据类型转换为 Numpy 数组的数据帧列会导致基于对象的 Numpy 数组，效率非常低（尤其是行主 Numpy 数组）。

请注意，除非输入数据帧已存储在 GPU 内存中，否则使用 GPU 不会使计算速度更快。事实上，CPU-GPU 数据传输比仅仅读取 RAM 更昂贵（由于互连开销通常是相当慢的 PCI 开销）。请注意，与 CPU 相比，GPU 内存非常有限。如果不需要传输目标数据帧，那么使用 cudf 相对简单，并且应该会带来较小的加速。为了获得更快的代码，需要实现快速的 CUDA 代码，但这对于具有混合数据类型的数据帧来说显然远非易事。最后，所产生的加速应该main_ram_throughput / gpu_ram_througput假设没有数据传输。注意这个系数一般是5-12。另请注意，CUDA 和 cudf 需要 Nvidia GPU。

最后，减少输入数据大小或仅减少计算量无疑是最佳解决方案（如@zvone 的评论中所示），因为它的计算量非常大。