调用scipy.ndimage.center_of_mass后内存暴增的原因及优化方法咨询

我正在使用scipy.ndimage模块获取3D数组中标记对象的相关信息，由于内存占用是限制计算服务器并行进程数量的瓶颈，因此尝试降低该过程的内存消耗。但发现调用center_of_mass函数后，内存占用急剧上升，且远大于函数返回数据的大小。

函数调用方式如下：

ndimage.center_of_mass(values, labels, range(1,np.max(labels)+1))

values和labels数组的形状均为(20, 2048, 2048)。调用center_of_mass前，我使用resource模块打印内存使用情况：

print(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss)

调用后再次打印，调用前数值为6504168，调用后为11255620，内存占用增加了近5GB。但center_of_mass返回数组的形状为(24187884, 3)，经检查ndarray的nbytes确认，该数组仅需约580MB内存。同时我用top监控进程，发现内存占用持续处于高位，并非函数调用期间的临时峰值。

请问导致内存大幅增长的原因是什么？有什么方法可以降低内存占用？

• 批量处理的临时数组开销：scipy.ndimage.center_of_mass一次性处理大量标签时，会为每个标签生成与原3D数组同尺寸的掩码或中间计算数组。当标签数量达到2400多万时，这些临时数组的总内存占用会远超过最终返回结果的大小；且部分临时内存可能因底层实现问题未被及时释放，导致内存持续高位。
• 底层内存回收不及时：scipy的ndimage模块部分实现基于C语言，若内部缓存或临时缓冲区未正确回收，即使Python层面的对象被删除，底层分配的内存也可能无法被系统回收，造成类似内存泄漏的占用情况。

• 分批处理标签：将所有标签拆分成小批次逐一处理，每处理完一批就手动清理临时变量并触发垃圾回收，避免一次性生成大量中间数组：

  import gc
  import numpy as np
  from scipy import ndimage
  
  max_label = np.max(labels)
  all_labels = range(1, max_label + 1)
  batch_size = 10000  # 根据服务器内存情况调整批次大小
  centers = []
  
  for idx in range(0, len(all_labels), batch_size):
      batch = all_labels[idx:idx+batch_size]
      batch_centers = ndimage.center_of_mass(values, labels, batch)
      centers.extend(batch_centers)
      # 手动清理临时变量
      del batch_centers
      gc.collect()
  
  final_centers = np.array(centers)
  

• 手动实现质心计算：利用numpy和ndimage的基础求和函数，直接计算各维度的加权和与总权重，避免scipy内部的批量处理开销：

  import numpy as np
  from scipy import ndimage
  
  # 获取所有有效标签（排除背景0）
  unique_labels = np.unique(labels[labels > 0])
  # 计算每个标签的总权重（values的和）
  total_weights = ndimage.sum(values, labels, unique_labels)
  # 生成各维度的坐标数组并计算加权和
  x_coords = np.arange(values.shape[0])[:, None, None]
  y_coords = np.arange(values.shape[1])[None, :, None]
  z_coords = np.arange(values.shape[2])[None, None, :]
  
  sum_x = ndimage.sum(values * x_coords, labels, unique_labels)
  sum_y = ndimage.sum(values * y_coords, labels, unique_labels)
  sum_z = ndimage.sum(values * z_coords, labels, unique_labels)
  # 计算质心
  final_centers = np.column_stack([sum_x/total_weights, sum_y/total_weights, sum_z/total_weights])
  

  这种方式仅生成几个与标签数量同大小的数组，内存占用远低于直接批量调用center_of_mass。
• 优化数组数据类型：若精度允许，将values从float64改为float32，labels改为uint32或更小的整数类型，减少原数组的内存占用，间接降低中间计算的内存开销。
• 提前计算最大标签值：将np.max(labels)的结果提前保存为变量，避免多次计算带来的额外内存和时间开销。

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