序列化共享内存的ndarray元组占用双倍空间，如何优化？

给定以下代码：

import numpy as np
import pickle

a = np.random.rand(4000)
b = a.T

assert np.shares_memory(a, b)

tup = (a, b)
s = pickle.dumps(tup)
print(len(s)) # 64187

s1 = pickle.dumps(a)
print(len(s1)) # 32151

由于a和b共享同一块内存，理论上序列化包含二者的元组时，结果大小应接近单独序列化其中一个的大小，但实际结果是两者大小之和。如何实现只序列化一份内存数据的目标？

解决方案 #

方法1：利用pickle备忘录机制手动复用引用 #

pickle的备忘录（memo）可以记录已序列化的对象，我们可以自定义序列化逻辑，让共享内存的数组视图复用已序列化的base数组，避免重复存储：

import numpy as np
import pickle

def pickle_shared_arrays(obj):
    memo = {}
    def reducer(item):
        if isinstance(item, np.ndarray):
            base = item.base if item.base is not None else item
            if base in memo:
                # 返回还原视图的逻辑，根据原数组生成转置/视图
                return (lambda arr: arr.T if item.flags.f_contiguous else arr, (memo[base],))
            memo[base] = item
        return pickle._reduce(item)
    
    return pickle.dumps(obj, protocol=pickle.HIGHEST_PROTOCOL, reducer=reducer)

# 测试
a = np.random.rand(4000)
b = a.T
tup = (a, b)
s = pickle_shared_arrays(tup)
print(len(s))  # 输出接近32151，与单独序列化a的大小一致

# 验证还原后的内存共享
loaded_tup = pickle.loads(s)
assert np.shares_memory(loaded_tup[0], loaded_tup[1])

方法2：使用numpy.savez（纯numpy数组场景） #

如果你的数据仅包含numpy数组，numpy.savez会自动检测共享内存的数组，只存储一份原始数据：

import numpy as np

a = np.random.rand(4000)
b = a.T
np.savez('shared_data.npz', a=a, b=b)

# 读取验证
loaded_data = np.load('shared_data.npz')
assert np.shares_memory(loaded_data['a'], loaded_data['b'])

生成的文件大小与单独存储a的大小接近。

方法3：修改numpy数组的pickle还原逻辑 #

通过全局修改np.ndarray.__reduce__方法，让数组视图只序列化base数组和视图参数，还原时重新生成视图：

import numpy as np
import pickle

# 保存原始reduce方法
original_reduce = np.ndarray.__reduce__

def shared_reduce(arr):
    if arr.base is not None:
        # 对于视图，返回生成视图的构造参数
        return (np.ndarray, (arr.shape, arr.dtype, arr.base, arr.offset, arr.strides))
    return original_reduce(arr)

np.ndarray.__reduce__ = shared_reduce

# 测试
a = np.random.rand(4000)
b = a.T
tup = (a, b)
s = pickle.dumps(tup, protocol=pickle.HIGHEST_PROTOCOL)
print(len(s))  # 输出接近32151

# 验证还原结果
loaded_tup = pickle.loads(s)
assert np.shares_memory(loaded_tup[0], loaded_tup[1])

# 恢复原始方法，避免影响其他逻辑
np.ndarray.__reduce__ = original_reduce

注意：该方法会全局修改numpy数组的序列化行为，需谨慎使用。

内容的提问来源于stack exchange，提问作者吴慈霆