joblib：选择 batch_size 和 pre_dispatch 来优化内存使用和性能答案

【问题标题】：joblib: Choosing batch_size and pre_dispatch to optimize memory usage and performancejoblib：选择 batch_size 和 pre_dispatch 来优化内存使用和性能
【发布时间】：2022-07-08 15:18:26
【问题描述】：

对于上下文，我正在尝试在频谱图列表上使用动态时间扭曲计算成对距离矩阵。所有声音文件都有不同的长度（时间步长），但在开始之前我知道它们的大小。该脚本按顺序运行良好，但计算时间太长，因此我尝试将其与 joblib 并行化。

假设我用 np.float32 类型的数组列表来表示它们（我将把所有代码放在下面的最小示例中）。作为替代，我将定义 dtw 函数来创建一个随机矩阵并返回最后一个单元格（行和列）中的值。我使用 numba 对其进行了优化，因此它运行得相当快。

import numpy as np
from joblib import Parallel, delayed


# Number of samples
n = 20000

# Generate
x = [np.random.uniform(size=(n, 40)) for n in np.random.randint(low=50, high=500, size=n)]

# Placeholder function
def fake_dtw(a, b):
    mat = np.random.uniform(size=(len(a), len(b)))
    return mat[-1, -1]

# Code to compute pairwise distance
batch_size = 1000
pre_dispatch = 2 * batch_size
with Parallel(n_jobs=-1, batch_size=batch_size, pre_dispatch=pre_dispatch) as p:
    results = p(
        delayed(
            lambda i, j, a, b: (i, j, fake_dtw(a, b))
        )(i, j, x[i], x[j])
        for i in range(1, len(x))
        for j in range(i)
    )

dtw_matrix = np.zeros(shape=(len(x), len(x)))
for i, j, res in results:
    dtw_matrix[i, j] = res
    dtw_matrix[j, i] = res

我已经阅读了documentation 以及这个问题What batch_size and pre_dispatch in joblib exactly mean。所以我知道 batch_size 和 pre_dispatch 是如何工作的，但我想不出一种方法来计算适当的值以获得最佳性能。

我的问题如下：给定

列表中所有项目的大小（我可以在启动前计算）
操作数（在本例中为 4 亿，因为它是 20000 个样本中的所有对）
CPU 数量（我一次最多可以启动 48 个工作器）
我电脑的 RAM (64 GB) 有没有办法让我选择batch_size 和pre_dispatch 以便尽可能快地计算运算？

在一个大约是我当前数据集大小 1/4 的数据集上，我能够使用 pre_dispatch='all' 和 batch_size=(number of operations)/os.cpu_count() 侥幸逃脱，因此所有数据在运行前会立即分发，但如果我尝试使用它会崩溃当前数据集（我认为这是由于内存使用造成的）。我尝试了更多值，但我想知道是否有更原则性的方法来执行此操作，而不是暴力破解并查看有效的方法。

提前谢谢你！

【问题讨论】：

标签： python parallel-processing joblib

【解决方案1】：

我从来没有完全找到问题本身的答案。我确实找到了解决方法；虽然我不确定它是否得到了最大程度的优化，但即使在较小的数据上运行速度也比以前快得多，并且在大约 4 小时内完成了在 20k 数据集上的运行。

所以我想我把这个解决方法留在这里以备将来参考，基本上是使用joblib documentation。

所以事实证明，我正在将 x 复制到由 joblib 生成的每个进程中，这在大约 20k 频谱图的情况下并不能很好地结束，更不用说。

根据文档，有几种解决方法：

使用带有threading 的线程而不是带有loky 的默认进程（我没有探索它，因为我无法使其与我的实现一起使用）
使用numpy memmap 来存储x 以及输出矩阵，然后定义它以便所有进程都可以访问它。这就是我最终使用并且运行良好的方法，在 4 个多小时内计算了约 2 亿个 DTW 距离（我承认我不完全确定它与 Python 中的其他实现相比如何，更不用说 C，但没有一个我发现的实现解决了我的问题，要么不适用于多元时间序列，要么不适用于不同时间长度的数组列表）

如果有人遇到这个问题，我将原始问题中的代码留在这里，更新后使用 memmap 解决方案：

import numpy as np
import os
import math
from joblib import Parallel, delayed


# Number of samples
n = 20000

# Generate
x = [np.random.uniform(size=(_, 40)) for _ in np.random.randint(low=50, high=500, size=n)]

# Placeholder function
def fake_dtw(i, j, a, b, output):
    mat = np.random.uniform(size=(len(a), len(b)))[-1, -1]
    output[i, j] = res
    output[j, i] = res

# Dump x to a memmap
memmap_folder = "temp"
x_filename_memmap = f'{memmap_folder}/x_memmap'
dump(x, x_filename_memmap)
x = load(x_filename_memmap, mmap_mode='r')

# Initialise the output matrix as writable memmap accessible by all processes
output_filename_memmap = f'{memmap_folder}/output_memmap'
output = np.memmap(output_filename_memmap, dtype=x[0].dtype, shape=(len(x), len(x)), mode='w+')

# Code to compute pairwise distance
batch_size = math.ceil(n*(n-1)/2 / os.cpu_count())
pre_dispatch = 'all'
with Parallel(n_jobs=-1, batch_size=batch_size, pre_dispatch=pre_dispatch) as p:
    results = p(
        delayed(
            lambda i, j, a, b: (i, j, fake_dtw(a, b))
        )(i, j, x[i], x[j], output)
        for i in range(1, len(x))
        for j in range(i)
    )

# retrieve output as you wish

【讨论】：