通过布尔掩码进行高效的 numpy 值分配

Question

我有一个布尔掩码值分配问题，需要有效的布尔掩码操作。

这是一个多维蒙版，我正在使用einsum来实现结果，但是操作效率不是很高，我想知道是否可以得到一些帮助 这是我目前的解决方案：（mask, truth_value, false_value 都是 dtype 和 shape 与我的问题匹配的虚拟数据。

mask = np.random.randn(1000, 50)> 0.5
truth_value = np.random.randn(50, 10)
false_value = np.random.randn(10)
objective = np.einsum('ij,jk->ijk', mask, truth_value) + np.einsum('ij,k->ijk', ~mask, false_value)

在给定mask, truth_value, false_value 的情况下，有没有更快的方法来获得objective？

在等待的时候，想出了一个更快的方法

objective = np.where(mask[...,np.newaxis], np.broadcast_to(truth_value, (1000, 50, 10)), np.broadcast_to(false_value,  (1000, 50, 10)))

但是有没有更快的选择？

Answer 1

您可以使用 Numba JIT 更有效地执行此操作。

import numpy as np
import numba as nb

@nb.njit('float64[:,:,::1](bool_[:,::1], float64[:,::1], float64[::1])')
def blend(mask, truth_value, false_value):
    n, m = mask.shape
    l = false_value.shape[0]
    assert truth_value.shape == (m, l)
    result = np.empty((n, m, l), dtype=np.float64)
    for i in range(n):
        for j in range(m):
            if mask[i, j]:
                result[i, j, :] = truth_value[j, :]
            else:
                result[i, j, :] = false_value[:]
    return result

mask = np.random.randn(1000, 50) > 0.5
truth_value = np.random.randn(50, 10)
false_value = np.random.randn(10)
objective = blend(mask, truth_value, false_value)

objective 在我的机器上的计算速度快了 4.8 倍。

如果这还不够快，您可以尝试使用参数 parallel=True 并在基于 i 的循环中使用 nb.prange 而不是 range 来并行化代码。由于创建新线程的开销，这可能不会更快。在我的机器（6 核）上，并行版本的速度快 7.4 倍（与执行时间相比，线程的创建非常昂贵）。

另一种可能的优化是将结果直接写入提前分配的缓冲区中（只有在使用相同数组大小多次调用此函数时才会更好）。

以下是我机器上的总体时序：

np.einsum:         4.32 ms
np.where:          1.72 ms
numba sequential:  0.89 ms
numba parallel:    0.58 ms