numba 的序列化类

Question

serializing/deserializing (pickling) numba 类有好的模式吗？以下是

    other = pkl.loads(pkl.dumps(jc))
  File "/usr/lib/python2.7/copy_reg.py", line 71, in _reduce_ex
    state = base(self)
TypeError: function takes at most 0 arguments (1 given)

我想 jit 很多类。

@jitclass([
    ('uncal_minRange', float64),
    ('uncal_maxRange', float64),
    ('nSize', int64),
    ('cal_percent', int64),
])
class SomeJitClass(object):
    def __init__(self, uncal_minRange=0.0, uncal_maxRange=0.0, nSize=0, cal_percent=0.0 ):
        self.uncal_minRange   = uncal_minRange
        self.uncal_maxRange   = uncal_maxRange
        self.cal_percent      = cal_percent
        self.nSize = nSize

    def compare(self, other):
        return self.uncal_minRange == other.uncal_minRange and \
            self.uncal_maxRange == other.uncal_maxRange and \
            self.cal_percent == other.cal_percent and \
            self.nSize == other.nSize


import unittest
class FastOpsTest(unittest.TestCase):

    def test_jit_class(self):
        import cPickle as pkl
        jc = SomeJitClass(1,2,3,4)
        print jc
        self.assertTrue(jc.compare(SomeJitClass(1,2,3,4)))
        self.assertFalse(jc.compare(SomeJitClass(1,2,3,5)))
        other = pkl.loads(pkl.dumps(jc))
        self.assertTrue(jc.compare(other))

Answer 1

下面是一个可能的解决方案，但顺便说一句，我发现很难推荐在其当前（2018 年初）实现中使用 jitclass - 只是有点过于限制/不成熟。如果可能的话，留在 numba-land 我会尝试使用普通数组/自由函数，and/or namedtuples 可以是捆绑一些相关状态的一个很好的解决方案。另外/或者 cython 是开发更丰富的类型的好工具，可以更好地与 python.

一起玩

也就是说，这是一个可能的解决方法。我正在使用一些 numba 内部机制，因此未来可能不稳定。

def dumps_jitclass(jc):
    typ = jc._numba_type_
    fields = typ.struct

    data = {
        'name': typ.classname,
        'struct': {k: getattr(jc, k) for k in fields}
    }
    return pickle.dumps(data)

def loads_jitclass(s):
    data = pickle.loads(s)
    cls = globals()[data['name']]
    instance = cls(**data['struct'])
    return instance

# usage
In [148]: jc = loads_jitclass(dumps_jitclass(SomeJitClass(1, 2, 3, 4)))

In [149]: oth = SomeJitClass(1, 2, 3, 4)

In [150]: jc.compare(oth)
Out[150]: True

Answer 2

感谢@chrisb 的回答。我为我的项目做了一些更改（来自装饰初学者）。

import pickle
from functools import wraps, update_wrapper

import numpy as np
from numba import jitclass, int32, float64

class jitpickler:
    '''
    pickler
    '''
    def __init__(self, jitobj):
        self.__dict__['obj'] = jitobj
        self.__dict__['__module__'] = jitobj.__module__
        self.__dict__['__doc__'] = jitobj.__doc__

    def __getstate__(self):
        obj = self.__dict__['obj']
        typ = obj._numba_type_
        fields = typ.struct

        return typ.classname, {k: getattr(obj, k) for k in fields}

    def __setstate__(self, state):
        name, value = state
        cls = globals()[name]
        value['_decorator'] = False
        jitobj = cls(**value)
        self.__init__(jitobj)

    def __getattr__(self, attr):
        return getattr(self.__dict__['obj'], attr)

    def __setattr__(self, attr, value):
        return setattr(self.__dict__['obj'], attr, value)

    def __delattr__(self, attr):
        return delattr(self.__dict__['obj'], attr)

def jitpickle(cls):
    decoratorkw = '_decorator'
    @wraps(cls)
    def decorator(*args, **kwargs):
        if kwargs.get(decoratorkw, True):
            kwargs.pop(decoratorkw, None)
            return jitpickler(cls(*args, **kwargs))
        else:
            kwargs.pop(decoratorkw, None)
            return cls(*args, **kwargs)
    return decorator

什么时候用

spec = [
    ('value', int32),
    ('array', float64[:])]
@jitpickle
@jitclass(spec)
class cls:
    def __init__(self, value, array):
        self.value = value
        self.array = array


a = cls(10, np.array([1.1,2.2]))
s = pickle.dumps(a)
b = pickle.loads(s)
print(b.value, b.array)

Ps：update_wrapper 不能很好地重载 __setattr__。而且我不想让 __setattr__ 太胖。

numba 的序列化 类

Serialization of numba classes

python

numba

numba 的序列化类