如何使用 numba.jit 方法

Question

在 python 中使用 numba.jit。

我可以将普通函数转换为 jit 类型运行:

from numba import jit

def sum(a, b):
    return a+b

func = jit(sum)
print(func(1, 2))

如何对方法执行此操作？像这样的东西（这不起作用，我知道为什么）。

from numba import jit

class some_class:
    def __init__(self, something = 0):
        self.number = something
    def get_num(self):
        return self.number

my_object = some_class(5)
func = jit(my_object.get_num)
print(my_object.func())

P.S。我也尝试过装饰器，它可以工作，但我不能将它用于导入的类（我自己没有定义的那些），所以我正在研究这个。

Answer 1

你不能 jit bound 方法，但你可以 jit 未绑定的方法（但只能在对象模式下）：

from numba import jit

class some_class:
    def __init__(self, something = 0):
        self.number = something
    def get_num(self):
        return self.number
    func = jit(get_num)

my_object = some_class(5)
print(my_object.func())
# 5

请注意，这不使用 nopython 模式，因此您不应期望任何合理的加速。您可以将 class 本身设为 jitclass（这意味着所有方法都不会自动进行 nopython-jitted），但它需要您键入属性：

import numba as nb

spec = [
    ('number', nb.int64),
]

@nb.jitclass(spec)
class some_class:
    def __init__(self, something):
        self.number = something
    def get_num(self):
        return self.number

my_object = some_class(5)
print(my_object.get_num())

但是对于更复杂的 classes，使用 jitclass 将变得非常困难（或不可能）。根据我的经验，最好的方法是简单地从方法中调用 jitted 函数：

from numba import njit  # like jit but enforces nopython-mode!

@njit
def my_func(val):
    return val  # this example is a bit stupid, I hope your real code does more!

class some_class:
    def __init__(self, something = 0):
        self.number = something
    def get_num(self):
        return my_func(self.number)

my_object = some_class(5)
print(my_object.get_num())

这取决于你的classand/or你的方法有多复杂，应该使用哪种方法。在你的情况下，我根本不会使用 numba，因为没有足够的计算昂贵的东西来补偿 numba 和 jit 开销。如果它有点复杂，我会使用 jitclass，如果它更复杂，我会使用从函数方法中调用的 jitted-function。就我个人而言，我永远不会将 jit 用于方法，因为这隐含地需要对象模式，因此 jitted 函数很可能比 unjitted 函数慢。

顺便说一句：在 Python 中，您通常使用 property 而不是 get_* 或 set_* 函数...

如何使用 numba.jit 方法

How to use numba.jit with methods

python

jit

class

function-parameter

numba