如何准确测量Python函数运行过程消耗的总内存？

两个现有方案失效的核心原因如下：

• 基于psutil读取RSS的方案统计的是操作系统视角下进程实际占用的物理内存差值，完全受Python内存池机制影响：Python通过pymalloc实现了多层内存池，对象销毁后内存不会立刻归还给操作系统，而是留在池中供后续复用。第一次执行函数时申请的大块内存，在del x后只是回到了Python内部内存池，没有还给OS，后续调用函数时直接从池中取内存，不会触发新的物理内存申请，因此RSS差值为0。同时这个方案完全统计不到中途申请又释放的内存，根本无法反映执行流程的总内存分配量。
• tracemalloc并非只能统计当前内存和峰值内存，多数人是没有用对它的能力：它本身可以跟踪Python层所有内存分配事件，完全可以实现类似Julia @btime的累计总内存统计效果。

方案1：原生tracemalloc实现（无额外依赖，统计Python层分配） #

适用于纯Python代码的内存统计，准确性高，不受内存池影响，多次调用结果稳定。Python3.12及以上版本可直接使用如下装饰器：

import tracemalloc
from functools import wraps

def profile_total_memory(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        # 兼容全局已启动tracemalloc的场景
        need_stop = not tracemalloc.is_tracing()
        if need_stop:
            tracemalloc.start()
        
        # 重置历史跟踪记录
        tracemalloc.clear_traces()
        total_alloc = 0

        # 注册分配回调，每次内存分配就累加大小，和后续是否释放无关
        def count_alloc(size, _):
            nonlocal total_alloc
            total_alloc += size
        tracemalloc.add_trace_callback(count_alloc)

        try:
            res = func(*args, **kwargs)
        finally:
            tracemalloc.remove_trace_callback(count_alloc)
            if need_stop:
                tracemalloc.stop()

        # 单位格式化
        def fmt(size):
            for u in ["B", "KB", "MB", "GB"]:
                if size < 1024 or u == "GB":
                    return f"{size:.2f} {u}"
                size /= 1024

        print(f"[{func.__name__}] 累计总分配内存: {fmt(total_alloc)}")
        cur_mem, peak_mem = tracemalloc.get_traced_memory()
        print(f"[{func.__name__}] 执行后存活内存: {fmt(cur_mem)}, 峰值内存: {fmt(peak_mem)}")
        return res
    return wrapper

使用时直接装饰目标函数即可，多次调用统计结果一致：

@profile_total_memory
def func():
    x = [1] * (10 ** 7)
    y = [2] * (4 * 10 ** 8)
    del x
    return y

# 连续调用两次
func()
func()

这个方案统计的是函数执行期间所有Python对象的内存申请总和，不管内存中途是否释放、是否走内存池复用，完全对应@btime的总内存统计逻辑。

如果你使用Python3.12以下版本，tracemalloc未开放分配回调接口，可以通过跟踪所有内存块的分配地址、记录新增分配大小的方式实现类似逻辑，只是代码复杂度稍高。

方案2：全链路内存跟踪（覆盖C扩展分配） #

如果函数中调用了numpy、pandas这类基于C实现的第三方库，C代码直接通过系统malloc申请的内存无法被tracemalloc捕获，此时可以使用挂钩libc内存分配接口的工具做统计。这类工具直接在操作系统层面拦截所有malloc/free调用，不会受Python内存池、C层内存池的影响，可以统计到函数执行期间所有的内存申请总量，包括中途释放的部分，统计结果最准确。

不要用RSS差值做精确内存统计：RSS是进程持有的物理页计数，受操作系统页回收策略、运行时内存池、GC时机的影响极大，只能做粗略的进程内存占用观测，完全无法用来统计单段代码的内存分配总量。

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