如何根据并行 Python 中的某些条件进行采样？

Question

假设我想根据条件并行采样。

例如，给出矩阵A。我想对p 对索引(i,j) 进行采样，这样A[i][j] != 5

import numpy as np
import random

A = np.random.randint(10, size=(5000, 5000)) # assume this is fixed

p = 400  # sample 400 index

res = set()
cnt = 0

while cnt < p:
    r, c = random.randint(0, A.shape[0]-1), random.randint(0, A.shape[0]-1)
    if A[r, c] != 5 and (r,c) not in res:
        res.add((r,c))
        cnt += 1

以上是我的尝试。但是，矩阵A 和样本数量p 可能非常大。我们可以并行吗？喜欢使用joblib、multiprocessing？或者有什么快速获取row和col的方法？

Answer 1

您可以使用 Numba 来加速此代码。 Numba 可以使用即时编译器 (JIT) 在运行时生成快速（并行）函数。使用像np.int8 这样的较小数据类型可以节省一些内存空间并缩短执行时间。实际上，较小的数组可以更快地从 RAM 读取/写入 RAM。此外，它们更有可能适合 CPU 缓存，从而加速随机访问。虽然您可以并行化随机选取，但这非常困难，并且线程的创建可能比有关所选参数的实际计算更昂贵。尽管如此，Numba 可以通过（大部分）消除 Python 解释器的开销来大幅提高其速度。 这是生成的代码：

# Initial conditions

import numba as nb
import numpy as np
import random

@nb.njit('int8[:,:](int_, int_)', parallel=True)
def genArray(n, m):
    res = np.empty((n, m), dtype=np.int8)
    # Parallel loop
    for i in nb.prange(n):
        for j in range(m):
            res[i, j] = np.random.randint(10)
    return res

p = 400
A = genArray(5000, 5000)

# Actual computing code

@nb.njit('(int8[:,::1], int_)')
def genPosSet(A, p):
    maxi = A.shape[0]-1
    res = set()
    cnt = 0
    while cnt < p:
        r, c = random.randint(0, maxi), random.randint(0, maxi)
        if A[r, c] != 5 and (r,c) not in res:
            res.add((r,c))
            cnt += 1
    return res

res = genPosSet(A, p)

genPosSet 的这个实现在我的机器上需要 64 us，而初始函数需要 1350 us。因此，新的实现速度快了 21 倍。

请注意，创建/删除线程（1 个线程/核心）并在它们之间共享工作的时间通常需要 10 到 1000 微秒。

请注意，如果p 不比A.size * prob 小很多，其中prob 是找到一个不同于5 的值的概率，那么当前的算法效率不是很高。在这种情况下，最好在选择随机位置之前过滤不同于 5 的值。如果p 不比A.size 小很多，那么最好的解决方案是打乱所有可以选择的位置，最后提取结果列表的p 第一个值。