Python机器学习应用内存占用远超tracemalloc检测值，如何排查？

问题背景 #

本地运行Python机器学习应用时，内存占用逐渐增长至700MB。内存峰值时，通过自定义代码统计全局/局部变量内存：

import sys
import pandas as pd
import numpy as np
global_vars = list(globals().items())
local_vars = list(locals().items())
all_vars = global_vars + local_vars
total = 0
mem = []
vars = []
for var, obj in all_vars:
    if isinstance(obj, pd.DataFrame):
        gso = obj.memory_usage().sum()
    elif isinstance(obj, np.ndarray):
        gso = obj.size
    else:
        gso = sys.getsizeof(obj)
    total += gso
    mem = mem + [gso]
    vars = vars + [var]
    print(var, gso)

df = pd.DataFrame({'obj': vars, 'size': mem}).sort_values('size', ascending=False)
print(df['size'].sum()/1000000)

统计结果仅略高于20MB。使用tracemalloc检测：

snapshot = tracemalloc.take_snapshot()
stats = snapshot.statistics('traceback')
sz = []
for stat in stats:
    sz = sz + [stat.size]
print(sum(sz)/1000000)

结果仅略低于100MB，与实际700MB的内存占用差距巨大，需要找出剩余内存的来源。

排查方向与解决方案 #

1. 修正Python对象内存统计的错误 #

你的自定义统计代码存在关键错误：

• 对于np.ndarray，用obj.size统计的是元素个数，而非实际内存占用，应替换为obj.nbytes（每个元素的字节数×元素个数）。
• sys.getsizeof仅返回对象本身的内存，不包含其引用的子对象（比如列表里的元素、字典的键值对），无法准确统计复杂对象的总内存。

2. 追踪C扩展/底层库的内存占用 #

Python原生工具无法统计C扩展（如NumPy、PyTorch、TensorFlow等）分配的原生内存，这是ML应用内存占用的主要来源之一：

• NumPy/Pandas：用df.memory_usage(deep=True).sum()统计DataFrame的真实内存（包含对象类型列的子对象）；用arr.nbytes统计数组内存。
• PyTorch/TensorFlow：检查是否有未释放的张量，尤其是被全局变量、闭包、生成器持有的张量。可以用torch.cuda.memory_summary()查看GPU内存细节，或手动清理张量引用后调用gc.collect()。
• 使用pympler库追踪全量对象内存：

  from pympler import muppy, summary
  all_objects = muppy.get_objects()
  summary.print_(summary.summarize(all_objects))
  

3. 检查进程外内存占用 #

• 子进程内存：如果应用使用multiprocessing或subprocess，子进程的内存不会被主进程的Python工具统计。用系统工具（Linuxhtop、Windows任务管理器）查看单个进程的内存占用，确认是否是子进程导致的高内存。
• 系统缓存/内存映射文件：操作系统会缓存大文件内容，这部分内存属于系统缓存，不算Python进程的常驻内存，但会被系统级工具统计。用free -h（Linux）查看buff/cache部分，或在Windows中查看“已修改”内存。

4. 优化tracemalloc的使用 #

• 尽早启动tracemalloc（脚本开头加入tracemalloc.start()），否则无法追踪启动后的内存分配。
• 对比不同时间点的快照，定位内存增长的代码段：

  import tracemalloc
  
  tracemalloc.start()
  snapshot1 = tracemalloc.take_snapshot()
  
  # 运行应用的核心逻辑
  snapshot2 = tracemalloc.take_snapshot()
  
  # 打印内存增长最多的10个调用栈
  stats = snapshot2.compare_to(snapshot1, 'lineno')
  for stat in stats[:10]:
      print(stat)
  

5. 排查循环引用与垃圾回收问题 #

• 手动触发垃圾回收：import gc; gc.collect()，再重新统计内存，看是否有下降。
• 查看无法自动回收的对象：

  import gc
  gc.set_debug(gc.DEBUG_SAVEALL)
  gc.collect()
  print([type(obj) for obj in gc.garbage])
  

  如果gc.garbage不为空，说明这些对象存在循环引用且包含无法自动释放的C扩展内存，需要手动清理引用。

6. 使用系统级工具分析内存 #

• Linux/macOS：用valgrind配合PYTHONMALLOC=malloc禁用Python内存池，追踪原生内存泄漏：

  PYTHONMALLOC=malloc valgrind --leak-check=full python your_script.py
  

• Windows：用VMMap工具查看进程的内存分布，区分私有内存、共享内存、内存映射文件等区域，定位占用最大的内存块。

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