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

Question

我想测量代码中每个for循环的RAM使用情况。我在互联网上搜索并找到了process = psutil.Process(os.getpid())和print(process.memory_info().rss)来测量RAM。但是这段代码获取整个进程的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

我很感激能够帮助我很多的回答。

- Orca

---

读取各部分之间的RAM使用情况，然后计算差异如何？ - maciek97x

@maciek97x 如果能够仅测量该块的 RAM 那将非常好，因为如果超过某个阈值，我想要中断该循环。 - Orca

@maciek97x，你能帮我写一些代码吗？我真的不知道该怎么做。 - Orca

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

正如您所看到的，在执行任何操作之前，该循环不会使用任何内存。

- maciek97x