python - 如何测量python中每个代码部分的RAM使用情况？

Question

我想测量代码中每个 for 循环的 RAM 使用情况。我在互联网上搜索并找到用于测量 RAM 的 process = psutil.Process(os.getpid()) 和 print(process.memory_info().rss)。但是这段代码获取了整个过程的 pid 而不是特定部分。有什么方法可以测量代码每个部分的 RAM 使用情况吗？例如，在下面的代码中，我们有 3 个 for 循环，它们填充了 3 个不同的字典。我想打印每个 for 循环和每个循环处理之间的 RAM 使用情况，如果 RAM 超过阈值，我想打破那个 for 循环。

dict1 = {}
dict2 = {}
dict3 = {}

for i in range (200):
  do something with dict1
  if RAM usage of this block exceeds 1GB then break
  this loop used: x Mb

for i in range (500):
  do something with dict2
  if RAM usage of this block exceeds 1GB then break
  this loop used: x2 Mb

for i in range (800):
  do something with dict3
  if RAM usage of this block exceeds 1GB then break
  this loop used: x3 Mb

我很感激能帮到我很多的答案

Answer 1

您可以在循环之前读取内存使用情况，然后在循环内再次读取它。然后您可以将循环内存使用量计算为这两个值之间的差值。如果超过某个阈值，则打破循环。

示例代码如下:

import numpy as np
import psutil
import os

process = psutil.Process(os.getpid())
a = []
threshhold = 64*1024*1024

base_memory_usage = process.memory_info().rss

for i in range(10):
    memory_usage = process.memory_info().rss
    loop_memory_usage = memory_usage - base_memory_usage
    
    print(loop_memory_usage)
    
    if loop_memory_usage > threshhold:
        print('exceeded threshold')
        break

    a.append(np.random.random((1000, 1000)))

结果:

0
8028160
16031744
24035328
32038912
40042496
48046080
56049664
64053248
72056832
exceeded threshold

如您所见，在执行任何操作之前，循环使用 0 字节的内存。