使用dill序列化/反序列化后numpy数组base属性指向无效内存而非None

问题原因 #

dill对不同大小的numpy数组采用差异化序列化策略：小数组会完整保留原数组的base属性，大数组则将数据转为bytes对象存储，反序列化时直接从该bytes重建数组，导致base指向bytes而非原有的None或数组自身，破坏了通过base判断原数组/视图的逻辑。

解决方案 #

通过自定义dill的序列化/反序列化钩子，手动记录数组的base状态，反序列化时重建符合预期的数组结构，同时兼容ndarray派生类及包含该类的对象。

1. 实现自定义序列化/反序列化逻辑

import numpy as np
import dill

def serialize_ndarray(obj):
    # 记录数组核心属性、base状态及对象类型
    return {
        'data': obj.tobytes(),
        'dtype': obj.dtype,
        'shape': obj.shape,
        'strides': obj.strides,
        'base_is_self': obj.base is obj,
        'base_is_none': obj.base is None,
        'type': type(obj)
    }

def deserialize_ndarray(state):
    # 从bytes重建基础数组
    arr = np.frombuffer(state['data'], dtype=state['dtype']).reshape(state['shape'])
    
    # 修正数组strides，保持原数组内存布局
    if state['strides'] != arr.strides:
        arr = arr.view()
        arr.strides = state['strides']
    
    # 根据记录的base状态修正数组
    if state['base_is_none']:
        # 复制数组使base变为None
        arr = arr.copy()
    elif state['base_is_self']:
        # 构造base指向自身的数组，利用buffer参数
        arr = np.ndarray(shape=arr.shape, dtype=arr.dtype, buffer=arr)
    
    # 转换为原ndarray派生类，如果不是原生ndarray
    if state['type'] is not np.ndarray:
        arr = arr.view(state['type'])
    
    return arr

2. 注册钩子到dill

将自定义逻辑注册为numpy数组的默认序列化/反序列化方式：

# 注册钩子，覆盖ndarray及子类的默认处理
dill.register(np.ndarray, serialize_ndarray, deserialize_ndarray)

3. 测试验证

# 测试普通数组
a = np.array(np.random.rand(100))
b = np.array(np.random.rand(200))
# 测试base为自身的数组
c = np.ndarray(shape=(5,), dtype=np.float64, buffer=np.array([1,2,3,4,5]))
# 测试ndarray派生类
class MyArray(np.ndarray):
    pass
d = np.array([1,2,3]).view(MyArray)

# 序列化反序列化
ad = dill.loads(dill.dumps(a))
bd = dill.loads(dill.dumps(b))
cd = dill.loads(dill.dumps(c))
dd = dill.loads(dill.dumps(d))

# 验证结果
print(f"a base: {type(a.base)}, ad base: {type(ad.base)}")
print(f"b base: {type(b.base)}, bd base: {type(bd.base)}")
print(f"c base: {type(c.base)}, cd base: {type(cd.base)}, cd.base is cd: {cd.base is cd}")
print(f"d type: {type(d)}, dd type: {type(dd)}, dd base: {type(dd.base)}")

说明 #

• 该方案会递归处理包含numpy数组的对象，无需额外修改容器类的序列化逻辑。
• 对于base为其他数组的视图，若需保留原base关联，可扩展序列化逻辑记录base的引用，但需注意循环引用问题。

内容的提问来源于stack exchange，提问作者Joce