在大型 numpy 数组上应用聚合函数答案

【问题标题】：Applying aggregation function over large numpy array在大型 numpy 数组上应用聚合函数
【发布时间】：2021-10-25 17:52:31
【问题描述】：

我正在尝试计算大型 numpy 数组的平均值。最初，我尝试过：

data = (np.ones((10**6, 133))
        for _ in range(100))
np.stack(data).mean(axis=0)

但我得到了

numpy.core._exceptions.MemoryError: 无法为形状为 (100, 1000000, 133) 且数据类型为 float32 的数组分配 xxx GiB

在原始代码中，数据是更有意义的向量的生成器。

我考虑过使用 dask 进行这样的操作，希望它将我的数据拆分成由磁盘支持的块。

import dask.array as da
import numpy as np

data = (np.ones((10**6, 133)) for _ in range(100))
x = da.stack(da.from_array(arr, chunks="auto") for arr in data)
x = da.mean(x, axis=0)
y = x.compute()

但是，当我运行它时，进程以“Killed”终止。

如何在单台机器上解决这个问题？

【问题讨论】：

平均值为1。抛开玩笑，也许this question 会引起你的兴趣。
@Ivan 我缺少 :haha 图标。谢谢。
@dzieciou 您在堆叠阵列时遇到此错误？
@MSS 是的，在堆叠时。

标签： python arrays numpy dask

【解决方案1】：

你可以试试这个方法：

agg_sum = np.zeros((10**6, 133))
total = 100

for dt in data:
    agg_sum = agg_sum + dt
_mean = agg_sum/total

【讨论】：

很好的解决方案！比我的解决方案运行得更快，并且需要更少的磁盘空间。我要补充两件事：（1）del dt 以避免内存不足，（2）用一些total 整数替换len(data)：预先知道要平均的数组总数，但你不能@987654325 @over 生成器。
我也不知道单个数组的形状（dt）所以我初始化了agg_sum = 0 ，它也很好用。
@dzieciou 我知道我们不能在生成器上使用len。我只是想展示这种方法。

【解决方案2】：

我发现的另一种解决方案是将所有数组存储在磁盘支持的文件中，使用 numpy.memmap。

import numpy as np

total = 100
shape = (10 ** 6, 133)
c = np.memmap(
    "total.array", dtype="float64", mode="w+", shape=(total, *shape), order="C"
)
for idx, arr in enumerate(data):
    c[idx,:,:] = arr[:]
    del arr
    
c.mean(axis=0)

这里重要的是del arr 避免在垃圾收集器回收未使用的数组之前使用整个内存。

请注意，该解决方案需要大约 100GB 的磁盘空间，而 @MSS 的解决方案通过仅保留当前总和所需的空间要少得多。

【讨论】：