在 for 循环中加载 Python pickle 会变慢

Question

我有一个 25GB 的 numpy 数组字典。该词典如下所示：

• 668,956 个键值对。
• 键是字符串。示例键： "109c3708-3b0c-4868-a647-b9feb306c886_1"
• 这些值是 shape 200x23、 type的 numpy 数组float64

当我在循环中重复使用 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

笔记：

• 在循环处理期间，RAM 使用量会逐渐下降（我假设取消引用file变量中的先前数据），然后再次上升。随着时间的推移，卸载和装载零件的速度似乎都会减慢。令我惊讶的是，卸载部分的 RAM 下降得如此之慢。
• 它增加的总 RAM 使用量保持不变（看起来不存在内存泄漏）。
• 我尝试过在循环中包含del fileand ，但这并没有加快任何速度。gc.collect()
• 如果我更改return pickle.load(handle)为return handle.read()，则卸载时间始终为 0.45 秒，加载时间始终为 4.85 秒。
• 我在 Windows 上使用带有 SSD 存储的 Python 3.9.13 ( Python 3.9.13 | packaged by conda-forge | (main, May 27 2022, 16:51:29) [MSC v.1929 64 bit (AMD64)])。
• 我有 64GB 内存，但似乎没有达到最大极限。
• 我为什么要这样做？在 ML 模型的训练过程中，我有 10 个文件，每个文件大小为 25GB。我无法同时将它们全部放入内存中，因此必须在每个时期加载和卸载它们。

有任何想法吗？如果有一种具有相似读取速度并且不会遇到上述问题的替代方案（我不担心压缩），我也愿意放弃使用 pickle。

编辑：我已经针对不同大小的泡菜运行了上述加载和卸载循环。下面的结果显示了速度随时间的相对变化。对于 3 GB 以上的任何内容，卸载时间开始显着增加。