如何使用Python的memory_profiler库统计代码内存使用量

问题原因说明 #

• 第一段代码报错原因：memory_usage 方法当参数 max_usage=True 时，会直接返回峰值内存的浮点数值，而非采样列表，因此代码中尝试用 res[0][0] 取值会抛出类型错误。
• 第二段批量检测写法不可行：proc 参数需要传入可调用的函数对象，你存储的是字符串形式的函数调用表达式，无法被直接执行。

单函数内存统计正确写法 #

如果只需要获取函数运行的峰值内存占用，参考如下实现：

import memory_profiler as mp

def fun(n):
    tmp = []
    for i in range(n):
        tmp.extend(list(range(i*i)))
    return "XXXXX"

# proc参数格式为 (函数对象, 参数列表, 参数字典)
res = mp.memory_usage(proc=(fun, [100], {}), max_usage=True, retval=True)
max_mem = res[0]
func_return = res[1]
print(f"函数运行峰值内存：{max_mem} MiB")
print(f"函数返回值：{func_return}")

如果需要获取运行过程中的完整内存采样序列，把 max_usage 设为False即可，此时返回的第一个值就是采样值列表：

res = mp.memory_usage(proc=(fun, [100], {}), max_usage=False, retval=True)
mem_samples = res[0]
print(f"运行过程内存采样值：{mem_samples}")
print(f"峰值内存：{max(mem_samples)} MiB")

批量检测多个函数内存占用实现 #

直接把函数对象存入列表即可，不需要转成字符串：

import memory_profiler as mp

def fun1(n):
    tmp = [i for i in range(n*1000)]
    return "XXXXX"
def fun2(n):
    tmp = [str(i) for i in range(n*1000)]
    return "YYYYY"

# 直接存储函数对象和对应参数，方便批量调用
methods = [
    (fun1, [100]),
    (fun2, [100])
]

for func, args in methods:
    max_mem, ret_val = mp.memory_usage(proc=(func, args, {}), max_usage=True, retval=True)
    print(f"函数 {func.__name__} 运行峰值内存：{max_mem:.2f} MiB，返回值：{ret_val}")

逐行内存分析方案 #

如果需要查看函数内部每行代码的内存占用，直接使用 @profile 装饰器，不需要写额外的统计代码：

1. 在需要分析的函数上添加 @profile 装饰器
2. 终端运行命令 mprof run 你的脚本名.py
3. 运行结束后执行 mprof plot 即可生成内存占用折线图，也可以直接在终端看到逐行的内存统计结果。

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