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

Posted 2023-04-18

【中文标题】如何以内存有效的方式在python中将矩阵一分为二？

【英文标题】：How to split a matrix in two in python in a memory effecient way?

【发布时间】：2014-10-06 09:00:24

【问题描述】：

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

这是我要优化的代码：

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。我不明白为什么？如何优化我的程序，使其功能相同但内存成本更低？

谢谢你，

山姆。

【参考方案1】：

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

您拆分此数组的方式非常规则，因此只需对其进行切片：

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

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

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

final_matrix = tmp1 + 1j * tmp2

或者更节省内存：

final_matrix = 1j * tmp2
final_matrix += tmp1

如果你只对总和感兴趣，你也可以将实部和虚部分别相加，最后组合起来。

【讨论】：

谢谢！通过这些修改，我的内存使用量从 11G 增加到了 2G 左右，这正是我想要的！