numba-jitclass
参考文档：https://numba.pydata.org/numba-doc/latest/user/jitclass.html#
jitclass
对类进行装饰
import numba as nb
import numpy as np
from numba.experimental import jitclass
spec = [
("value", nb.int32),
("array", nb.float32[:]),
]
@jitclass(spec)
class Bag:
def __init__(self, value):
# self.value = value
self.array = np.zeros(value, dtype=np.float32)
@property
def size(self):
return self.array.size
def increment(self, val):
for i in range(self.size):
self.array[i] += val
return self.array
a = 2
b = Bag(a)
print("b.increment(1): ", b.increment(1)) # b.increment(1): [1. 1.]
上面例子中，spec 提供了一个两元组元素的数组，元组包含字段名称和类型。也可以使用有序字典映射字段与类选的关系。
类中至少要初始化每个定义的字段，如果不初始化，字段会包含垃圾数据。
具体的 numb.typed 容器(container)做类成员
- 显式的类型和构建
kv_ty = (nb.types.int64, nb.types.unicode_type)
@jitclass(
[("d", nb.types.DictType(*kv_ty)), ("l", nb.types.ListType(nb.types.float64))]
)
class ContainerHolder:
def __init__(self):
self.d = nb.typed.Dict.empty(*kv_ty)
self.l = nb.typed.List.empty_list(nb.types.float64)
c = ContainerHolder()
c.d[1] = "apple"
c.d[2] = "orange"
c.l.append(1.0)
c.l.append(2.0)
print("c.d: ", c.d) # c.d: {1: apple, 2: orange}
print("c.l: ", c.l) # c.l: [1.0, 2.0]
- 另一个有用的模式是用 numba.typed 的 _numba_type_ 属性能够找找到容器的类型, 这样可以直接在 python 解释器中访问容器的实例。使用 numba.typeof 可以得到跟容器实例一样的信息。如下：
d = nb.typed.Dict()
d[1] = "apple"
d[2] = "orange"
l = nb.typed.List()
l.append(1.0)
l.append(2.0)
@jitclass([("d", nb.typeof(d)), ("l", nb.typeof(l))])
class ContainerInsHolder:
def __init__(self, dict_instance, list_instance):
self.d = dict_instance
self.l = list_instance
c = ContainerInsHolder(d, l)
print("c.d: ", c.d) # c.d: {1: apple, 2: orange}
print("c.l: ", c.l) # c.l: [1.0, 2.0]
c.d[3] = "banana"
print("c.d: ", c.d) # c.d: {1: apple, 2: orange, 3: banana}
需要注意的是，容器实例在使用前必须要初始化，否则会有异常，如下面的是错误的：
d_ty = nb.types.DictType(nb.types.int64, nb.types.unicode_type)
@jitclass([("d", d_ty)])
class NotInitContainer:
def __init__(self):
self.d[10] = "apple" # d 没有被初始化，这里是无效的
NotInitContainer() # 实例化会失败，内存访问无效，程序会异常结束 Process finished with exit code -1073741819 (0xC0000005)
支持的操作
以下 jitclasses 操作在 python 解释器和 numba 编译的函数中都支持：
- 用 jitclass 类实例化对象。（如： bag = Bag(123)）
- 读/写属性。（如：bag.value）
- 方法调用。（如：bag.increment(2)）
- 调用实例的静态方法。（如：bag.add(1, 2)）
- 调用类的静态方法。（如：Bag.add(1,2)）
局限性
- jitclass 被看作是一个 numba 的编译函数
- isinstance() 只能在 python 解释器中使用
完整代码
import numba as nb
import numpy as np
from numba.experimental import jitclass
spec = [
("value", nb.int32),
("array", nb.float32[:]),
]
@jitclass(spec)
class Bag:
def __init__(self, value):
# self.value = value
self.array = np.zeros(value, dtype=np.float32)
@property
def size(self):
return self.array.size
def increment(self, val):
for i in range(self.size):
self.array[i] += val
return self.array
a = 2
b = Bag(a)
print("b.increment(1): ", b.increment(1)) # b.increment(1): [1. 1.]
print("-" * 20)
kv_ty = (nb.types.int64, nb.types.unicode_type)
@jitclass(
[("d", nb.types.DictType(*kv_ty)), ("l", nb.types.ListType(nb.types.float64))]
)
class ContainerHolder:
def __init__(self):
self.d = nb.typed.Dict.empty(*kv_ty)
self.l = nb.typed.List.empty_list(nb.types.float64)
c = ContainerHolder()
c.d[1] = "apple"
c.d[2] = "orange"
c.l.append(1.0)
c.l.append(2.0)
print("c.d: ", c.d) # c.d: {1: apple, 2: orange}
print("c.l: ", c.l) # c.l: [1.0, 2.0]
print("-" * 20)
d = nb.typed.Dict()
d[1] = "apple"
d[2] = "orange"
l = nb.typed.List()
l.append(1.0)
l.append(2.0)
@jitclass([("d", nb.typeof(d)), ("l", nb.typeof(l))])
class ContainerInsHolder:
def __init__(self, dict_instance, list_instance):
self.d = dict_instance
self.l = list_instance
c = ContainerInsHolder(d, l)
print("c.d: ", c.d) # c.d: {1: apple, 2: orange}
print("c.l: ", c.l) # c.l: [1.0, 2.0]
c.d[3] = "banana"
print("c.d: ", c.d) # c.d: {1: apple, 2: orange, 3: banana}
print("-" * 20)
d_ty = nb.types.DictType(nb.types.int64, nb.types.unicode_type)
@jitclass([("d", d_ty)])
class NotInitContainer:
def __init__(self):
self.d[10] = "apple" # d 没有被初始化，这里是无效的
NotInitContainer() # 实例化会失败，内存访问无效，程序会异常结束 Process finished with exit code -1073741819 (0xC0000005)
