Cython 中 numpy 数组掩码的性能

Posted 2023-02-23

【中文标题】Cython 中 numpy 数组掩码的性能

【英文标题】：Performance of numpy array masking in Cython

【发布时间】：2017-08-29 15:07:57

【问题描述】：

作为这个问题的后续here（感谢 MSeifert 的帮助），我想出了一个问题，即在将屏蔽数组传递给之前，我必须用索引数组 new_vals_idx 屏蔽一个 numpy 数组更新val_dict。

对于在旧帖子中回答 MSeifert 提出的解决方案，我尝试应用数组掩码，但性能并不令人满意。 我用于以下示例的数组和字典是：

import numpy as np
val_dict = 'a': 5.0, 'b': 18.8, 'c': -55/2
for i in range(200):
    val_dict[str(i)] = i
    val_dict[i] = i**2

keys = ('b', 123, '89', 'c')  # dict keys to update
new_values = np.arange(1, 51, 1) / 1.0  # array with new values which has to be masked
new_vals_idx = np.array((0, 3, 5, -1))  # masking array
valarr = np.zeros((new_vals_idx.shape[0]))  # preallocation for masked array
length = new_vals_idx.shape[0]

为了使我的 code-sn-ps 更容易与我的旧问题进行比较，我将坚持使用 MSeifert 答案的函数命名。这些是我从 python/cython 中获得最佳性能的尝试（由于性能太差，其他答案被忽略了）：

def old_for(val_dict, keys, new_values, new_vals_idx, length):
    for i in range(length):
        val_dict[keys[i]] = new_values[new_vals_idx[i]]
%timeit old_for(val_dict, keys, new_values, new_vals_idx, length)
# 1000000 loops, best of 3: 1.6 µs per loop

def old_for_w_valarr(val_dict, keys, new_values, valarr, new_vals_idx, length):
    valarr = new_values[new_vals_idx]
    for i in range(length):
        val_dict[keys[i]] = valarr[i]
%timeit old_for_w_valarr(val_dict, keys, new_values, valarr, new_vals_idx, length)
# 100000 loops, best of 3: 2.33 µs per loop

def new2_w_valarr(val_dict, keys, new_values, valarr, new_vals_idx, length):
    valarr = new_values[new_vals_idx].tolist()
    for key, val in zip(keys, valarr):
        val_dict[key] = val
%timeit new2_w_valarr(val_dict, keys, new_values, valarr, new_vals_idx, length)
# 100000 loops, best of 3: 2.01 µs per loop

Cython 函数：

%load_ext cython
%%cython
import numpy as np
cimport numpy as np
cpdef new3_cy(dict val_dict, tuple keys, double[:] new_values, int[:] new_vals_idx, Py_ssize_t length):
    cdef Py_ssize_t i
    cdef double val  # this gives about 10 µs speed boost compared to directly assigning it to val_dict
    for i in range(length):
        val = new_values[new_vals_idx[i]]
        val_dict[keys[i]] = val
%timeit new3_cy(val_dict, keys, new_values, new_vals_idx, length)
# 1000000 loops, best of 3: 1.38 µs per loop

cpdef new3_cy_mview(dict val_dict, tuple keys, double[:] new_values, int[:] new_vals_idx, Py_ssize_t length):
    cdef Py_ssize_t i
    cdef int[:] mview_idx = new_vals_idx
    cdef double [:] mview_vals = new_values
    for i in range(length):
        val_dict[keys[i]] = mview_vals[mview_idx[i]]
%timeit new3_cy_mview(val_dict, keys, new_values, new_vals_idx, length)
# 1000000 loops, best of 3: 1.38 µs per loop

# NOT WORKING:
cpdef new2_cy_mview(dict val_dict, tuple keys, double[:] new_values, int[:] new_vals_idx, Py_ssize_t length):
    cdef double [new_vals_idx] masked_vals = new_values
    for key, val in zip(keys, masked_vals.tolist()):
        val_dict[key] = val

cpdef new2_cy_mask(dict val_dict, tuple keys, double[:] new_values, valarr, int[:] new_vals_idx, Py_ssize_t length):
    valarr = new_values[new_vals_idx]
    for key, val in zip(keys, valarr.tolist()):
        val_dict[key] = val

Cython 函数 new3_cy 和 new3_cy_mview 似乎并不比 old_for 快很多。传递valarr 以避免函数内部的数组构造（因为它将被调用数百万次）甚至似乎减慢了它。 在 Cython 中使用 new_vals_idx 数组在 new2_cy_mask 中进行屏蔽会给我错误：“指定的内存视图的索引无效，请键入 int[:]”。对于索引数组，有没有像 Py_ssize_t 这样的类型？ 尝试在 new2_cy_mview 中创建屏蔽内存视图时出现错误“无法将类型 'double[:]' 分配给 'double [__pyx_v_new_vals_idx]'”。甚至有像蒙面记忆视图这样的东西吗？我无法找到有关此主题的信息...

将时间结果与我的旧问题的结果进行比较，我猜想数组屏蔽是占用大部分时间的过程。而且由于它很可能已经在 numpy 中进行了高度优化，因此可能没什么可做的。但是减速是如此巨大，必须（希望）有更好的方法来做到这一点。 任何帮助表示赞赏！提前致谢！

【问题讨论】：

我不确定你会比new3_cy 做得更好——这基本上是一个 python 重函数（从 tuple 读取并写入 dict），所以 15-20% 的加速范围是关于你通常能得到什么。

好的，感谢您提供的信息！我可以很容易地摆脱元组，并且只需一点​​点努力就可以摆脱字典。只是 numpy 数组需要 bei Usedom。您期望什么加速以及您推荐哪些数据类型？

如果您需要字符串和/或异构类型和/或哈希查找，很难击败内置的dict。如果您的问题以某种方式减少到仅使用单个数字类型和 numpy 数组，那么 10 倍以上的加速并不少见，这取决于。

好的，那么我将尝试引入数字 ID，而不是带字符串的 dictlookups。还有tuples？由于不变性对于我的代码的普通用户来说是一个很好的特性，并且由于纯 python 速度很高，所以我认为使用tuples 可能是个好主意。我应该改用 numpy 数组吗？我可以将它们设置为 array.flags.writeable = False，但是 afaik 使它们与 Cython 不兼容......好吧，不一定需要不变性......

这是一种权衡，如果您当前的方法足够快，我会继续使用 dict 和元组，更简单、更安全。但是为了最大限度地提高性能，是的，您基本上希望所有内容都在 numpy 数组或等效数组中，以便您可以在 c 级别处理数据。

【参考方案1】：

您可以在当前构造中做的一件事是关闭边界检查（如果安全的话！）。不会有很大的不同，但会增加一些性能。

%%cython
import numpy as np
cimport numpy as np
cimport cython

@cython.boundscheck(False)
@cython.wraparound(False)
cpdef new4_cy(dict val_dict, tuple keys, double[:] new_values, int[:] new_vals_idx, Py_ssize_t length):
    cdef Py_ssize_t i
    cdef double val  # this gives about 10 µs speed boost compared to directly assigning it to val_dict
    for i in range(length):
        val = new_values[new_vals_idx[i]]
        val_dict[keys[i]] = val

In [36]: %timeit new3_cy(val_dict, keys, new_values, new_vals_idx, length)
1.76 µs ± 209 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

In [37]: %timeit new4_cy(val_dict, keys, new_values, new_vals_idx, length)
1.45 µs ± 31.5 ns per loop (mean ± std. dev. of 7 runs, 1000000 loops each)

【讨论】：

我非常需要环绕。不使用它会降低我的程序对我以外的其他用户的模块化。只有 diabling boundscheck 是一个小的加速，但我想一旦我测试了所有东西，它会是值得的。