Python Numpy 广播 4D 阵列吃 RAM

Python Numpy broadcasting 4D array eating RAM

我正在尝试对 4D 图像数组(包括时间序列)进行简单计算,但与初始化数组相比,广播会占用大量 RAM,而且我已经尝试阅读其他有类似问题的文章.例如。

这里有一条来自“rth”的评论说 “广播不会为初始数组分配额外的内存”

来自“Warren Weckesser”的已接受评论显示提问者使用 og newaxis 创建一个额外的分配数组的问题。

我尝试按照 Warren 展示的方法进行操作,但我仍然吃掉了很多 RAM,我不明白为什么。现在我已经实现了 rths 块计算方法,结果很好,但它仍然困扰着我为什么直接 numpy 计算在 RAM 使用中爆炸。

这是我所做的一个例子

我初始化我将向其添加数据的数组,并创建 uint16 的随机原始图像,因为它来自具有 12 位 RAW 图像数据的 16 位 TIFF 文件。我保留其余的 float32 以节省 RAM。最后一个精度没那么重要

import numpy as np
imagearraydensity = np.ones((512, 1024, 250, 20),dtype=np.float32)
imagearrayraw = np.random.randint(4095,size=(512, 1024, 250, 20),dtype=np.uint16)

我有两个预先计算的线性常数数组,这里只是随机数

acons = np.random.random((512, 1024)).astype(np.float32)
bcons = np.random.random((512, 1024)).astype(np.float32)

然后计算

np.divide(np.subtract(imagearrayraw, bcons[:, :, np.newaxis, np.newaxis], dtype=np.float32),
                      acons[:, :, np.newaxis, np.newaxis],
                      imagearraydensity,  # Output array position
                      dtype=np.float32)

我的带有真实数据的代码最终在我的系统上使用了大约 36 GB RAM,在完成计算后使用了大约 27 GB。

我能做些什么来减少直接广播对 RAM 的使用,还是我已经实现的基于块的最佳方法?

首先,两个输入数组包含 2_621_440_000 项,产生 15 GiB 的 RAM。 np.subtract 生成的临时数组包含相同数量的元素,从而产生 10 GiB 的 RAM。输出数组被覆盖,所以它不应该占用更多的 RAM。这意味着 Numpy 应该已经占用了 25 GiB 的 RAM。

问题是 Numpy 不知道如何从 float32 数组中减去 uint16 数组。因此它应用 导致 uint16 数组被转换为 float32 数组。这意味着它创建了一个新的 10 GiB 临时数组,因此占用了大约 36 GiB 的 RAM。请注意,Numpy 和 Python 都不是为了最小化内存占用而设计的。

也就是说,有多种解决方案可以解决此问题。一种解决方案是将输入数据分成块,以便临时数组占用更少 space。这做起来相当麻烦艰难。另一个理论上的解决方案是自己转换数组,以便将其写入 imagearrayraw,但 Numpy 没有为 astype 提供 out 参数。另一种解决方案是使用 Numba 以便即时执行转换(效率更高)。

这是 chunk-based 解决方案:

for i in range(imagearraydensity.shape[0]):
    imagearraydensity[i,...] = imagearrayraw[i,...].astype(np.float32)
np.subtract(imagearraydensity, bcons[:, :, np.newaxis, np.newaxis], out=imagearraydensity, dtype=np.float32)
np.divide(imagearraydensity, acons[:, :, np.newaxis, np.newaxis], out=imagearraydensity, dtype=np.float32)

此解决方案在我的机器上占用 15 GiB。