gpt4 book ai didi

python - 如何使用 Numba 加速 python 的字典

转载 作者:行者123 更新时间:2023-12-01 15:49:41 25 4
gpt4 key购买 nike

我需要在 bool 值数组中存储一些单元格。起初我使用 numpy,但是当数组开始占用大量内存时,我有了一个想法,即在字典中以元组作为键存储非零元素(因为它是可散列类型)。例如:{(0, 0, 0): True, (1, 2, 3): True} (这是“3D 数组”中的两个单元格,索引为 0,0,0 和 1,2,3,但维数是预先未知的,并且在我运行算法时定义)。
它有很大帮助,因为非零单元只填充了数组的一小部分。

为了从这个字典中写入和获取值,我需要使用循环:

def fill_cells(indices, area_dict):
for i in indices:
area_dict[tuple(i)] = 1

def get_cells(indices, area_dict):
n = len(indices)
out = np.zeros(n, dtype=np.bool)
for i in range(n):
out[i] = tuple(indices[i]) in area_dict.keys()
return out

现在我需要用 Numba 加速它。 Numba 不支持原生 Python 的 dict(),所以我使用了 numba.typed.Dict。
问题是 Numba 在定义函数的阶段想知道键的大小,所以我什至无法创建字典(键的长度事先未知,并在调用函数时定义):
@njit
def make_dict(n):
out = {(0,)*n:True}
return out

Numba 无法正确推断字典键的类型并返回错误:
Compilation is falling back to object mode WITH looplifting enabled because Function "make_dict" failed type inference due to: Invalid use of Function(<built-in function mul>) with argument(s) of type(s): (tuple(int64 x 1), int64)

如果我在函数中将 n 更改为具体数字,它就可以工作。我用这个技巧解决了它:
n = 10
s = '@njit\ndef make_dict():\n\tout = {(0,)*%s:True}\n\treturn out' % n
exec(s)

但我认为这是错误的低效方式。而且我需要将我的 fill_cells 和 get_cells 函数与@njit 装饰器一起使用,但是 Numba 返回相同的错误,因为我试图在此函数中从 numpy 数组创建元组。

我了解 Numba 的基本限制(以及一般的编译),但也许有一些方法可以加快功能,或者,也许你对我的单元存储问题有另一种解决方案?

最佳答案

最终解决方案:

主要问题是 Numba 在定义创建元组的函数时需要知道元组的长度。诀窍是每次都重新定义功能。我需要使用定义功能的代码生成字符串并使用 运行它执行() :

n = 10
s = '@njit\ndef arr_to_tuple(a):\n\treturn (' + ''.join('a[%i],' % i for i in range(n)) + ')'
exec(s)

之后我可以调用 arr_to_tuple(a) 创建可以在另一个 @njit - 装饰函数中使用的元组。

例如,创建元组键的空字典,需要解决问题:
@njit
def make_empty_dict():
tpl = arr_to_tuple(np.array([0]*5))
out = {tpl:True}
del out[tpl]
return out

我在字典中写了一个元素,因为它是 Numba 推断类型的一种方式。

另外,我需要使用我的 填充单元格 get_cells 问题中描述的功能。这就是我用 Numba 重写它们的方式:

写作元素。只是将 tuple() 更改为 arr_to_tuple():
@njit
def fill_cells_nb(indices, area_dict):
for i in range(len(indices)):
area_dict[arr_to_tuple(indices[i])] = True

从字典中获取元素需要一些令人毛骨悚然的代码:
@njit
def get_cells_nb(indices, area_dict):
n = len(indices)
out = np.zeros(n, dtype=np.bool_)
for i in range(n):
new_len = len(area_dict)
tpl = arr_to_tuple(indices[i])
area_dict[tpl] = True
old_len = len(area_dict)
if new_len == old_len:
out[i] = True
else:
del area_dict[tpl]
return out

我的 Numba (0.46) 版本不支持 .contains (in) 运算符和 try-except 构造。如果您有支持它的版本,您可以为它编写更多“常规”解决方案。

因此,当我想检查字典中是否存在具有某个索引的元素时,我会记住它的长度,然后我在字典中写一些带有提到索引的东西。如果长度改变,我得出结论该元素不存在。否则元素存在。看起来很慢的解决方案,但事实并非如此。

速度测试:

解决方案的工作速度惊人。与 native Python 优化代码相比,我使用 %timeit 对其进行了测试:
  • arr_to_tuple() 比普通 快 5 倍元组() 功能
  • get_cells 与 numba 相比,一个元素快 3 倍,大元素数组快 40 倍 native Python 编写的 get_cells
  • 使用 numba 填充单元格 相比,一个元素快 4 倍,大元素数组快 40 倍 native Python 编写的 fill_cells
  • 关于python - 如何使用 Numba 加速 python 的字典,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/59490782/

    25 4 0
    Copyright 2021 - 2024 cfsdn All Rights Reserved 蜀ICP备2022000587号
    广告合作:1813099741@qq.com 6ren.com