时间戳

首页 > 解决方案 > 将 Numba jitclass 与来自字典和元组的输入一起使用

python - 将 Numba jitclass 与来自字典和元组的输入一起使用

问题描述

我正在优化我拥有的一些代码，这些代码主要包含在单个 python 类中。它对 python 对象的操作很少，所以我认为使用 Numba 将是一个很好的匹配，但我在创建对象期间需要大量参数，而且我认为我不完全理解 Numba 相对较新的 dict 支持（此处的文档）。我拥有的参数都是单个浮点数或整数，并被传递到对象中，存储，然后在整个代码运行过程中使用，如下所示：

import numpy as np
from numba import jitclass, float64

spec = [
    ('p', dict),
    ('shape', tuple),               # the shape of the array
    ('array', float64[:,:]),          # an array field
]

params_default = {
    par_1 = 1,
    par_2 = 0.5
    }

@jitclass(spec)
class myObj:
    def __init__(self,params = params_default,shape = (100,100)):
        self.p = params
        self.shape = shape
        self.array = self.p['par_2']*np.ones(shape)

    def inc_arr(self):
        self.array += self.p['par_1']*np.ones(shape)

有很多我认为我不明白 Numba 需要什么。如果我想使用 nopython 模式使用 Numba 优化它，是否需要将规范传递给 jitclass 装饰器？如何定义字典的规范？我还需要声明形状元组吗？我查看了我在 jitclass 装饰器上找到的文档以及 dict numba 文档，但我不知道该怎么做。当我运行上面的代码时，我收到以下错误：

TypeError: spec values should be Numba type instances, got <class 'dict'>

我是否需要以某种方式在规范中包含 dict 元素？从文档中不清楚正确的语法是什么。

或者，有没有办法让 Numba 推断输入类型？

标签： pythondictionaryjitnumba

解决方案

spec需要由特定于 numba 的类型组成，而不是由 python 类型组成！因此tuple，dict在规范中必须输入numba 类型（并且 afaik 只允许使用同质字典）。

因此，您可以在 jitted 函数中指定您的params_defaultdict，如此处所示，或者您显式键入一个 numba dict ，如此处所示。

对于这种情况，我将采用后一种方法：

import numpy as np
from numba import jitclass, float64

# Explicitly define the types of the key and value:
params_default = nb.typed.Dict.empty(
    key_type=nb.typeof('par_1'),
    value_type=nb.typeof(0.5)
)

# assign your default values
params_default['par_1'] = 1.  # Same type required, thus setting to float
params_default['par_2'] = .5

spec = [
    ('p', nb.typeof(params_default)),
    ('shape', nb.typeof((100, 100))),               # the shape of the array
    ('array', float64[:, :]),          # an array field
]

@jitclass(spec)
class myObj:
    def __init__(self, params=params_default, shape=(100, 100)):
        self.p = params
        self.shape = shape
        self.array = self.p['par_2'] * np.ones(shape)

    def inc_arr(self):
        self.array += self.p['par_1'] * np.ones(shape)

正如已经指出的那样： dict 是同质类型的。因此，所有键/值必须属于同一类型。所以存储int和float在同一个字典中是行不通的。

推荐阅读