循环中重复加载Python pickle速度变慢的原因及解决方案

问题背景 #

现有一个25GB的Pickle文件，存储结构为包含668,956个键值对的字典：

• 键为字符串，示例："109c3708-3b0c-4868-a647-b9feb306c886_1"
• 值为形状200x23、类型float64的NumPy数组

在循环中重复加载该文件时，加载和"卸载"（对象回收）速度随迭代次数逐渐变慢，具体表现见下方代码与运行结果。

测试代码 #

import time
import pickle

def load_pickle(file: int) -> dict:
    with open(f"D:/data/batched/{file}.pickle", "rb") as handle:
        return pickle.load(handle)

for i in range(0, 9):
    print(f"\nIteration {i}")
    
    start_time = time.time()
    file = None
    print(f"Unloaded file in {time.time() - start_time:.2f} seconds")

    start_time = time.time()
    file = load_pickle(0)
    print(f"Loaded file in {time.time() - start_time:.2f} seconds")

运行结果 #

Iteration 0
Unloaded file in 0.00 seconds
Loaded file in 18.80 seconds

Iteration 1
Unloaded file in 14.78 seconds
Loaded file in 30.51 seconds

Iteration 2
Unloaded file in 28.67 seconds
Loaded file in 30.21 seconds

Iteration 3
Unloaded file in 35.38 seconds
Loaded file in 40.25 seconds

Iteration 4
Unloaded file in 39.91 seconds
Loaded file in 41.24 seconds

Iteration 5
Unloaded file in 43.25 seconds
Loaded file in 45.57 seconds

Iteration 6
Unloaded file in 46.94 seconds
Loaded file in 48.19 seconds

Iteration 7
Unloaded file in 51.67 seconds
Loaded file in 51.32 seconds

Iteration 8
Unloaded file in 55.25 seconds
Loaded file in 56.11 seconds

补充现象与测试 #

• 循环中内存占用先降后升，卸载、加载速度随迭代持续变慢，内存下降速度超出预期
• 内存峰值稳定，无内存泄漏迹象
• 手动执行del file和gc.collect()无法改善性能
• 若改为直接返回文件二进制内容（return handle.read()），卸载时间稳定在0.45s，加载稳定在4.85s
• 环境：Windows + SSD，Python 3.9.13（conda-forge），64GB内存未耗尽
• 场景：机器学习训练时，每个epoch需加载/卸载10个25GB文件，无法同时存入内存
• 补充测试：文件体积超过3GB时，卸载时间会出现显著上升；加载时间也随文件体积增大持续上升

原因分析 #

1. Pickle反序列化的对象开销：Pickle加载时会逐个重建Python对象（字典、NumPy数组），反复创建/销毁数百万个对象会导致Python内存池碎片化，垃圾回收（GC）的开销随迭代次数累积——代码中file = None后的"卸载时间"实际是GC在回收前一次的大量对象，内存碎片化又会导致后续加载时内存分配速度变慢。
2. NumPy数组的序列化特性：Pickle对NumPy数组的序列化基于对象而非连续二进制块，反复重建大量数组会进一步加剧内存碎片化。
3. Windows内存管理特性：大内存块的频繁分配与释放可能触发系统级内存整理，多次重复后开销逐渐累积。

解决方案 #

替代Pickle的方案（推荐） #

以下方案均为二进制格式，读写效率稳定，无性能衰减问题：

1. NumPy原生格式（.npy/.npz）

利用NumPy的高效二进制存储，将字典拆分为键列表和大数组：

# 保存数据
import numpy as np

data = load_pickle(0)  # 调用原Pickle加载函数
keys = np.array(list(data.keys()), dtype=str)
# 将所有数组拼接为(N, 200, 23)的大数组
arrays = np.stack(list(data.values()), axis=0)

np.save("D:/data/batched/keys.npy", keys)
np.save("D:/data/batched/arrays.npy", arrays)

# 加载数据
def load_numpy_data():
    keys = np.load("D:/data/batched/keys.npy")
    arrays = np.load("D:/data/batched/arrays.npy")
    return dict(zip(keys, arrays))

2. HDF5格式（适合大规模结构化数据）

使用h5py库存储，支持高效随机读写，内存管理更稳定：

# 安装依赖：pip install h5py
import h5py

# 保存数据
data = load_pickle(0)
with h5py.File("D:/data/batched/data.h5", "w") as f:
    for key, arr in data.items():
        f.create_dataset(key, data=arr, dtype=np.float64)

# 加载数据
def load_hdf5_data():
    data = {}
    with h5py.File("D:/data/batched/data.h5", "r") as f:
        for key in f.keys():
            data[key] = f[key][:]
    return data

3. MsgPack配合MsgPack-NumPy

高效二进制序列化格式，支持NumPy数组，性能接近原生二进制：

# 安装依赖：pip install msgpack msgpack-numpy
import msgpack
import msgpack_numpy as m

# 启用NumPy序列化支持
msgpack_numpy.patch()

# 保存数据
data = load_pickle(0)
with open("D:/data/batched/data.msgpack", "wb") as f:
    msgpack.dump(data, f)

# 加载数据
def load_msgpack_data():
    with open("D:/data/batched/data.msgpack", "rb") as f:
        return msgpack.load(f)

Pickle格式的临时优化 #

若必须保留Pickle格式，可尝试：

• 使用pickle.HIGHEST_PROTOCOL重新序列化文件，减少体积与反序列化开销
• 将大文件拆分为多个小Pickle文件，分批加载以降低单次GC压力
• 显式导入import _pickle as pickle，确保使用最快的C实现版本

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