python - 用于节省内存的 Numpy nditer？

Question

在使用 nditer 遍历 ndarray 时我迷路了。

背景

我正在尝试为 3D 数组中的每个点计算 3x3 对称矩阵的特征值。 我的数据是一个形状为 [6,x,y,z] 的 4D 数组，其中 6 个值是矩阵在 x、y、z 点的值，在 ~500x500x500 的 float32 立方体上。 我首先使用了 numpy 的 eigvalsh，但它针对大型矩阵进行了优化，而我可以对 3x3 对称矩阵使用分析简化。

然后我实现了 wikipedia's simplification ，两者都是采用单个矩阵并计算特征值的函数(然后使用嵌套 for 循环简单地迭代)，然后使用 numpy 进行矢量化。

问题是，现在在我的矢量化中，每个操作都会创建一个我的数据大小的内部数组，最终导致使用过多的 RAM 和 PC 卡住。

我尝试使用 numexpr 等，它仍然在 10G 左右使用。

我想做什么

我想遍历(使用 numpy 的 nditer)我的数组，以便为​​每个矩阵计算我的特征值。这将消除分配巨大的中间数组的需要，因为我们一次只计算 ~ 10 个 float 。 基本上是尝试将嵌套的 for 循环替换为一个迭代器。

我正在寻找这样的东西:

for a,b,c,d,e,f in np.nditer([symMatrix,eigenOut]): # for each matrix in x,y,z

    # computing my output for this matrix
    eigenOut[...] = myLovelyEigenvalue(a,b,c,d,e,f)

到目前为止我最好的是:

for i in np.nditer([derived],[],[['readonly']],op_axes=[[1,2,3]]):

但这意味着 i 获取 4D 数组的所有值，而不是长度为 6 的元组。 我似乎无法理解 nditer 文档。

我做错了什么？关于遍历“除了一个”轴之外，您有什么提示和技巧吗？

关键是要有一个 nditer，它会在迭代时胜过常规的嵌套循环(一旦这有效，我将更改函数调用、缓冲区迭代……但到目前为止，我只希望它能正常工作 ^^)

Answer 1

你真的不需要 np.nditer 来做这个。遍历除第一个轴以外的所有轴的更简单方法是将其 reshape 为 [6, 500 ** 3] 数组，将其转置为 [500 ** 3, 6]，然后遍历行:

for (a, b, c, d, e, f) in (symMatrix.reshape(6, -1).T):
    # do something involving a, b, c, d, e, f...

如果你真的想使用 np.nditer 那么你会做这样的事情:

for (a, b, c, d, e, f) in np.nditer(x, flags=['external_loop'], order='F'):
    # do something involving a, b, c, d, e, f...

需要考虑的潜在重要事项是，如果 symMatrix 是 C 顺序(行优先)而不是 Fortran 顺序(列优先)，那么在第一个维度上迭代可能比迭代最后 3 个维度，从那时起您将访问相邻的内存地址 block 。因此，您可能要考虑切换到 Fortran 顺序。

我不希望从其中任何一个中获得巨大的性能提升，因为在一天结束时，您仍然在 Python 中执行所有循环并且仅在标量上运行，而不是利用矢量化。