用 Python tracemalloc 追踪内存分配：实用教程与示例

在使用 Python 编程时，你通常无需过多关注内存分配的细节。但追踪内存分配有时却非常有用，尤其是在处理内存密集型操作和大型数据集时。

Python 内置的 tracemalloc 模块提供了一些函数，可以帮助你了解内存的使用情况并调试应用程序。借助 tracemalloc，你可以获知内存分配的位置与数量，拍摄内存快照，对比快照差异等等。

本教程将介绍其中的一些用法。让我们开始吧！

开始之前

我们将使用一个简单的 Python 脚本进行数据处理。为此，我们会创建一个示例数据集并对其进行处理。除了较新版本的 Python，你还需要在工作环境中安装 pandas 和 NumPy。

创建并激活一个虚拟环境：

$ python3 -m venv v1

$ source v1/bin/activate

安装所需的库：

$ pip3 install numpy pandas

你可以在 GitHub 上找到本教程的代码。

创建包含订单详情的示例数据集

我们将生成一个包含订单详情的示例 CSV 文件。运行以下脚本，即可创建一个含有 10 万条订单记录的 CSV 文件：

# create_data.py

import pandas as pd

import numpy as np

# 创建包含订单详情的示例数据集

num_orders = 100000

data = {

  'OrderID': np.arange(1, num_orders + 1),

  'CustomerID': np.random.randint(1000, 5000, num_orders),

  'OrderAmount': np.random.uniform(10.0, 1000.0, num_orders).round(2),

  'OrderDate': pd.date_range(start='2023-01-01', periods=num_orders, freq='min')

}

df = pd.DataFrame(data)

df.to_csv('order_data.csv', index=False)

此脚本会生成一个 pandas DataFrame，含有 10 万条记录，并将其导出为 CSV 文件。包含以下四个字段：

OrderID

：每个订单的唯一标识符

CustomerID

：客户编号

OrderAmount

：订单金额

OrderDate

：订单的日期和时间

使用 tracemalloc 追踪内存分配

接下来，我们将编写一个 Python 脚本来加载并处理数据集，并追踪内存分配情况。

首先，定义load_data和process_data两个函数，分别用于从 CSV 文件加载数据和处理数据记录：

# main.py

import pandas as pd

def load_data(file_path):

  print("Loading data...")

  df = pd.read_csv(file_path)

  return df

def process_data(df):

  print("Processing data...")

  df['DiscountedAmount'] = df['OrderAmount'] * 0.9  # 应用 10% 折扣

  df['OrderYear'] = pd.to_datetime(df['OrderDate']).dt.year  # 提取订单年份

  return df

然后，我们按照以下步骤追踪内存分配：

使用tracemalloc.start()初始化内存追踪。

调用load_data()读取 CSV 文件为 DataFrame。在这一步后拍摄一次内存快照。

调用process_data()，为 DataFrame 添加两个新列：DiscountedAmount和OrderYear。处理后再拍摄一次快照。

对比两次快照，找出内存使用差异，并打印出内存消耗最多的代码行。

打印当前和峰值内存用量，以评估整体影响。

示例代码如下：

import tracemalloc

def main():

  # 开始追踪内存分配

  tracemalloc.start()

  # 加载数据

  df = load_data('order_data.csv')

  # 拍摄第一次快照

  snapshot1 = tracemalloc.take_snapshot()

  # 处理数据

  df = process_data(df)

  # 拍摄第二次快照

  snapshot2 = tracemalloc.take_snapshot()

  # 对比快照

  top_stats = snapshot2.compare_to(snapshot1, 'lineno')

  print("[ 内存消耗最多的代码行 ]")

  for stat in top_stats[:10]:

      print(stat)

  # 当前与峰值内存用量

  current, peak = tracemalloc.get_traced_memory()

  print(f"当前内存用量: {current / 1024 / 1024:.1f} MB")

  print(f"峰值用量: {peak / 1024 / 1024:.1f} MB")

  tracemalloc.stop()

if __name__ == "__main__":

  main()

运行脚本：

$ python3 main.py

输出示例：

Loading data...

Processing data...

[ 内存消耗最多的代码行 ]

/home/balapriya/trace_malloc/v1/lib/python3.11/site-packages/pandas/core/frame.py:12683: size=1172 KiB (+1172 KiB), count=4 (+4), average=293 KiB

/home/balapriya/trace_malloc/v1/lib/python3.11/site-packages/pandas/core/arrays/datetimelike.py:2354: size=781 KiB (+781 KiB), count=3 (+3), average=260 KiB

:123: size=34.6 KiB (+15.3 KiB), count=399 (+180), average=89 B

当前内存用量: 10.8 MB

峰值用量: 13.6 MB

总结

使用 tracemalloc 追踪内存分配可以帮助你识别内存消耗较大的操作，并借助返回的内存追踪数据和统计信息来优化程序性能。

你可以据此判断是否需要采用更高效的数据结构和处理方法来减少内存使用。对于长时间运行的应用程序，也可以定期使用 tracemalloc 监控内存用量。此外，你还可以将 tracemalloc 与其他分析工具结合使用，全面了解内存使用情况。

如果你想了解如何用 memory-profiler 进行内存分析，请阅读《Python 内存分析入门》。