Python函数性能优化求助：进程崩溃的数组运算函数优化

崩溃根源 #

你的my_function崩溃的核心原因是生成了规模为n×n的巨型矩阵：当输入data的长度为n时，reshaped_data / reshaped_data.transpose()会通过广播生成n×n的数组，内存占用呈平方级增长。比如n=20万时，float64类型的矩阵会占用320GB内存，直接超出你的可用内存（215.45GB），触发内存不足导致进程崩溃；同时全量计算也会浪费CPU资源。

优化方案 #

根据你的硬件（64逻辑核、大内存），分两种场景给出优化：

场景1：不需要完整的n×n矩阵（优先推荐）

如果后续仅需访问矩阵的部分元素/行/列，完全没必要生成全量矩阵，改用按需计算：

import numpy as np

# 获取单个元素[i,j]的比值
def get_ratio(data, i, j):
    arr = np.asarray(data).ravel()
    return arr[i] / arr[j]

# 获取第i行的所有比值
def get_row_ratios(data, i):
    arr = np.asarray(data).ravel()
    return arr[i] / arr

这种方式内存占用仅为O(n)，不会触发崩溃，计算效率也最高。

场景2：必须生成完整的n×n矩阵

如果业务逻辑要求必须保留全量矩阵，从内存优化和并行加速两方面改进：

1. 减少内存开销

   • 用ravel()替代flatten()：flatten()会生成数组副本，ravel()返回原数组的视图（输入为连续数组时），节省内存和拷贝时间；
   • 降低数据精度：如果业务允许，将数组转为float32，内存占用直接减半；
   • 简化广播代码：避免显式reshape，直接用索引实现广播，代码更简洁高效。

   优化后的基础版本：

   import numpy as np
   
   def my_function_optimized(data):
       arr = np.asarray(data).ravel()
       # 直接利用numpy广播生成矩阵，省略多余的reshape步骤
       result = arr[:, np.newaxis] / arr[np.newaxis, :]
       # 可选：转为float32节省内存
       # result = result.astype(np.float32)
       return result
   

2. 利用多CPU并行加速
   结合你的64逻辑核，用numba的JIT并行编译加速行级计算，充分利用CPU资源：

   import numpy as np
   from numba import njit, prange
   
   @njit(parallel=True, fastmath=True)
   def my_function_parallel(data):
       arr = np.asarray(data).ravel()
       n = len(arr)
       # 预先分配内存，避免动态扩容开销
       result = np.empty((n, n), dtype=arr.dtype)
       # 并行遍历每一行，独立计算行内所有比值
       for i in prange(n):
           result[i] = arr[i] / arr
       return result
   

   这个版本会自动利用所有逻辑核，计算速度比纯numpy版本提升数倍，同时内存控制更高效。

额外建议 #

在调用函数前，可以先判断输入规模：如果n过大（比如n>15万时，float64矩阵会超过200GB），直接抛出警告或切换到按需计算模式，避免崩溃。

内容的提问来源于stack exchange，提问作者Beck