Python追踪内存占用

DechinPhy

发布于 2023-12-21 08:55:55

2.5K0

技术背景

当我们需要对python代码所占用的内存进行管理时，首先就需要有一个工具可以对当前的内存占用情况进行一个追踪。虽然在Top界面或者一些异步的工具中也能够看到实时的内存变化，还有一些工具可以统计代码中每一步的内存占用。但如果只是要查看单步操作之后的内存变化，tracemalloc的简单易用，让它成为了一个绝佳的选择。本文主要介绍用tracemalloc来追踪代码的内存占用变化。

tracemalloc的使用

tracemalloc的操作逻辑非常简单，在开始统计时使用一个start函数，结束统计的时候使用一个stop函数，中间过程就像拍照片一样不断的使用get_traced_memory函数即可。

import numpy as np
import tracemalloc

a = np.random.random((1000000,))
tracemalloc.start()
current, peak = tracemalloc.get_traced_memory()
print(f"Current memory usage is {current / 10**6}MB; Peak was {peak / 10**6}MB")
b = np.random.random((1000000,))
current, peak = tracemalloc.get_traced_memory()
print(f"Current memory usage is {current / 10**6}MB; Peak was {peak / 10**6}MB")
b = b.astype(np.float32)
current, peak = tracemalloc.get_traced_memory()
print(f"Current memory usage is {current / 10**6}MB; Peak was {peak / 10**6}MB")
del a
current, peak = tracemalloc.get_traced_memory()
print(f"Current memory usage is {current / 10**6}MB; Peak was {peak / 10**6}MB")
del b
current, peak = tracemalloc.get_traced_memory()
print(f"Current memory usage is {current / 10**6}MB; Peak was {peak / 10**6}MB")
tracemalloc.stop()

在上面这个案例中，我们在统计内存占用前也分配了一个numpy数组，但是我们发现这个内存分配被自动忽略了。也就是说，我们只统计start函数开始之后的每一步的操作导致的内存变化。我们在start之后定义了一个numpy数组b，这里还是一个numpy.float64的数组，占用了8MB的内存。在将其转化为numpy.float32的数组之后，内存一下子缩小了一半，为4MB。但是我们发现，此时的峰值内存占用是12MB，也就是说，这个astype的操作，其实相当于定义了一个新的数组，然后把原数组拷贝到新的数组中，再将原数组释放掉这样的一个流程。在测试案例的最后，我们使用python的del删除了这个数组对象，此时内存就被清空了。具体运行结果如下所示：

Current memory usage is 0.0MB; Peak was 0.0MB
Current memory usage is 8.000096MB; Peak was 8.000312MB
Current memory usage is 4.000152MB; Peak was 12.000248MB
Current memory usage is 4.000152MB; Peak was 12.000248MB
Current memory usage is 5.6e-05MB; Peak was 12.000248MB