python - 如何以内存有效的方式在Python中将矩阵一分为二？

Question

我有一个很大的多维矩阵，其中一个索引实际上包含数字的实部和虚部。

这是我想要优化的代码:

import numpy as np

big_matrix = np.random.random((8,160,23,3,23,80)) # 1240M

tmp1 = np.zeros((8,80,23,3,23,80)) # 620M
tmp2 = np.zeros((8,80,23,3,23,80)) # 620M

for ii in np.arange(80):
  tmp1[:,ii,:,:,:,:] = big_matrix[:,2*ii,:,:,:,:]
  tmp2[:,ii,:,:,:,:] = big_matrix[:,2*ii+1,:,:,:,:]

final_matrix = np.vectorize(complex)(tmp1,tmp2) # 1240M

a = np.sum(final_matrix)

big_matrix 的理论内存大小应为 (8*160*23*3*23*80)*8/(1024**2)=1240MB。所以我预计总内存消耗为 3.7 GB。相反，我的内存消耗上升到了 11GB。我不懂为什么？如何优化我的程序，使其执行相同的操作但内存成本更便宜？

谢谢，

萨姆.

Answer 1

据我了解，numpy.vectorize本质上是一个Python循环，因此效率非常低。您看到的高内存消耗很可能是由它引起的。

分割这个数组的方式非常规则，所以只需对其进行切片即可:

tmp1 = big_matrix[:,  ::2, ...]
tmp2 = big_matrix[:, 1::2, ...]

这会创建原始数组的“ View ”，因此不需要额外的内存。

查看答案 here ，构造复杂数组的简单方法是:

final_matrix = tmp1 + 1j * tmp2

或者更高的内存效率:

final_matrix = 1j * tmp2
final_matrix += tmp1

如果您只对总和感兴趣，您也可以分别对实部和虚部进行求和，最后将它们合并起来。