用于节省内存的 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 个浮点数。基本上是尝试将嵌套的 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 文档。

我做错了什么？关于遍历 "all but one" 轴，您有什么提示和技巧吗？

关键是要有一个 nditer，它会在迭代时胜过常规的嵌套循环（一旦这个工作我会改变函数调用，缓冲迭代......但到目前为止我只是想让它工作^^）

Answer 1

你真的不需要 np.nditer。遍历除第一个轴以外的所有轴的更简单方法是将其重新整形为 [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 个维度上，从那时起您将访问相邻的内存地址块。因此，您可能要考虑切换到 Fortran 顺序。

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

用于节省内存的 Numpy nditer？

Numpy nditer for memory saving?

python

arrays

iteration

numpy