在非常大的 numpy 数组上进行计算：流计算与核外内存

计算科学 Python 麻木的 内存管理
2021-12-27 08:08:02

我正在尝试在 numpy 中执行取决于多个参数的计算，并涉及创建许多中间数组。这些中间步骤涉及更多参数的积分。正因为如此，在我计算的部分过程中，我尝试创建一个数组，我意识到它的低端大小约为 30GB。我的系统无法处理这个。

作为我遇到的问题的一般示例，请考虑以下 Python 代码：

param1s = np.linspace(min1, max1, 100)
param2s = np.linspace(min2, max2, 100)
param3s = np.linspace(min3, max3, 100)

def main(param1, param2, param3):
    dummy_param_as = np.linspace(0, 10, 100)
    dummy_param_bs = np.linspace(-1, 1, 100)

    huge_array = do_big_calculation(params1, params2, params3,
                                    dummy_param_as, dummy_param_bs)

    integral_over_dummy_b = np.trapz(huge_array, x=dummy_param_bs axis=-1)

    return np.trapz(integral_over_dummy_b, x=dummy_param_as, axis=-1)

具体来说，do_big_calculation将所有参数沿单独的轴一起广播，创建一个 shape 数组(100, 100, 100, 100, 100)。然而，这个数组不需要存在很长时间，直到它可以集成到它的最后两个轴上。

对于这种情况，似乎有一个强烈推荐的解决方案：使用h5py或dask之类的东西将数据写入存储，并通过从存储的文件中加载块中的数据来执行计算。这种类型的解决方案似乎通常被称为核外方法。

但是，我发现了另一种似乎很少被讨论的途径：将结果流出来并迭代地构造结果，而无需触及存储。我发现的一个这样的解决方案是npstreams。

对于我的玩具示例，这种方法的工作方式是huge_array在内存中逐步构建梯形规则总和，以便huge_array根本不会加载到内存中。npstreams看起来像是一种实现方式。

我很好奇，考虑到我正在计算我的数据而不是加载存储的数据，这些方法的优缺点是什么？一个可能会更快吗？一个或两个都存在缩放问题？

1个回答

其中很多将取决于您的do_big_calculation功能的细节。

通常，出于性能原因，您希望避免将数据推送到磁盘。磁盘 I/O 速度明显低于内存速度。

有一些策略可能有助于避免首先创建那个巨大的矩阵。

如果每个位置的输出值仅取决于参数，即您可以编写

def do_big_calculation(a,b,c,d,e):
   result = np.zeros((len(a),len(b),len(c),len(d),len(e)):
   for ia in range(a):
       for ib in range(b):
           for ic in range(c):
               for id in range(d):
                   for ie in range(e):
                       result[ia,ib,ic,id]=do_big_calculation_core(
                             a[ia],b[ib],c[ic],d[id],e[ie]
                       )
   return result

那么创建大矩阵没有任何优势，甚至不需要流式传输，您只需在循环内进行梯形积分即可。

def do_big_calculation_and_integrate(a,b,c,d,e):
   result = np.zeros((len(a),len(b),len(c)):
   for ia in range(a):
       for ib in range(b):
           for ic in range(c):
               temp1 = np.zeros(len(d))
               for id in range(d):
                   temp2 = np.zeros(len(e))
                   for ie in range(e):
                       temp2[ie]=do_big_calculation_core(
                             a[ia],b[ib],c[ic],d[id],e[ie]
                       )
                   temp1[id] = np.trapz(temp2,...)
               result[ia,ib,ic] = np.trapz(temp1)
   return result

当然，这种显式循环会破坏 numpy 中的性能，因此您需要小心 - 在这种情况下，重新排序循环和组合会有所帮助（以稍微增加内存为代价。我还发现 numba包可以相当好地优化这些循环（特别是如果您可以将 do_big_calculation_core 编写为 numba 函数）。

另一种选择是生成切片或块，并将它们累积到最终结果中。如果参数值之间共享大量工作，或者 numpy+numba 在显式循环中阻塞，您可能需要执行此操作。

其它你可能感兴趣的问题
上一篇病态矩阵系统的迭代解 下一篇在 Rayleigh-Bernard 模拟中包含拉普拉斯项的指数技巧是什么