Python 性能分析的几个方法，找到你代码中的那个她！

我们在编写了一个脚本在笔记本上处理一些数据，然后去喝杯咖啡或者上了个厕所，15分钟后回来时发现进度才完成不到10%。

我们的脑袋里面就会发问：为什么这么慢？究竟是在哪个部分是慢的？是读取数据、处理数据还是保存数据？如何让它变快？它真的很慢吗?

有了这个疑问我们尝试去解决这个问题，下面我们介绍几个 python 性能分析的工具。

什么是性能分析

分析器是一种工具，它接收代码，运行它，并收集关于每个函数调用花费了多少时间、被执行了多少次以及函数调用的层次结构等信息。

通过分析其输出，我们可以找出代码中哪部分执行时间最长（称为瓶颈），甚至可能想出解决方法。希望识别和解决瓶颈问题，从而实现整体速度上的最大改进。

问题示例：

假设我们有一个大文本文件，想要在其中找到某个特定模式的多个出现。首先，让我们生成一个由随机字母数字字符组成的大文件（逐行生成）：

 import random
 import string
 
 def generate_random_string(length):
     """生成一个由小写字母和数字组成的随机字符串。"""
     letters_and_digits = string.ascii_letters + string.digits
     return ''.join(random.choice(letters_and_digits) for _ in range(length))
 ​
 def generate_random_file(filename, num_lines, line_length):
     """生成一个指定数量和长度的随机行文件。"""
     with open(filename, 'w') as f:
         for i in range(num_lines):
             random_line = generate_random_string(line_length) + '\n'
             f.write(random_line)
 ​
 if __name__ == '__main__':
     #生成1000000行，每行1000个字符
     generate_random_file('random.txt', 1_000_000, 1000)  

接下来，让我们定义基准函数——它逐行读取文件，然后计算以数字为前缀的单词bob和rob在其中出现的次数：

 import re
 ​
 def baseline():
     num_total_matches = 0
     pattern1 = r"[0-9]{1}bob" # 正则表达式
     pattern2 = r"[0-9]{1}rob" # 正则表达式
     
     with open('random.txt', 'rt') as f:
         for line in f:
             line = line.lower()  # 将文本全部为小写。
             for pattern in [pattern1, pattern2]:
                 # 在字符串中找到模式的所有匹配项
                 matches = re.findall(pattern, line)
 ​
                 # 统计出现的次数
                 num_matches = len(matches)
                 num_total_matches += num_matches
             
     return num_total_matches

例如，一行abc1robdef02bob中有两个出现。

我们运行baseline函数，计算出现次数（在我的情况下为10861），并测量运行时间 — 在我的机器上为32秒。

我们如何加快速度？

改进目标

代码中有四个潜在部分导致运行速度变慢。

1. 使用原始字符串而不是编译后的正则表达式对象。
2. 使用两个单独的搜索，而不是一个联合的正则表达式。
3. 将每一行转换为小写，而不是使用不区分大小写标志。
4. 逐行读取文件，而不是较大的块。

我们如何知道哪一个是瓶颈？我们可以使用性能分析工具。

Python性能分析工具

我们不必自己实现任何功能代码。Python已经有两个内置的性能分析模块 — cProfile 和 profile。它们执行相同的功能，但 cProfile 是用C编写的，而 profile 则是纯Python编写的。你可以直接使用它们，不过我们使用一些外部工具，这将大大提升我们的工作效率。

我们需要一种快速便捷的方式来对代码的一部分（例如一个函数）进行性能分析，并将结果保存到文件中。profilehooks模块提供了一个简单装饰器，我们可以像这样使用到我们需要检测的函数：

 from profilehooks import profile
 ​
 # stdout=False -> don't print anything in the terminal
 # filname -> path to the output file with profiling results
 @profile(stdout=False, filename='baseline.prof')
 def baseline():
     ...

可以通过简单的 pip install profilehooks 命令进行安装。

我们需要以人类可读的方式来展示这个文件。我使用以下两种工具进行操作。

SnakeViz — 简单快速

Snakeviz 是 Python 分析结果的交互式可视化工具，可以在浏览器中使用。

安装非常简单（pip install snakeviz），使用也很方便（snakeviz <path-to-profiling-output>）。

让我们看看我们的 baseline 函数的分析结果如何。

Snakeviz 交互式结果 — icicle（左侧）和sunburst（右侧，可能不太易读）图表。我们可以悬停并单击每个函数调用以查看其详细信息。从上到下是调用的嵌套层次结构，线条的长度表示代码执行该部分所花费的相对时间。

Snakeviz 还显示了一个交互式表格，其中包含函数执行时间的统计信息。

我们可以看到大部分执行时间都花费在findall函数内部，也就是进行正则表达式搜索。这意味着，如果我们想加快代码运行速度，需要专注于加快该函数的速度，因为它是瓶颈所在，并非代码的其他部分。

让我们用第二个工具来确认结果。

gprof2dot — 可读性强且灵活

Gprof2dot 提供更易读的可视化，以流程图的形式保存为图片文件，因此很容易分享（并在必要时自动化）结果。然而它不是交互式的，并且需要系统中安装Graphviz。

要安装 gprof2dot，只需使用 pip install gprof2dot 命令。

要生成带有性能分析结果的输出图像，请使用以下命令：

 python -m gprof2dot -f pstats <profiling-results-file> | dot -Tpng -o output.png

我们首先将函数调用的层次结构表示为dot格式的图形，然后生成一幅图像 — 显示该图形。dot命令支持不同的输出格式，包括.jpg和.svg，而且 gprof2dot 的输出也具有高度可定制性。

我们代码中函数执行层次结构的图形表示。百分比显示了在函数内部花费的总执行时间比例. 但这些仅包括自身函数的代码。最后一个值是该函数在代码中被调用的次数。

可以说现在更易读了，图表告诉我们正则表达式搜索占据了总执行时间的88%，因此我们需要专注于让它变得更快。

如果无法访问Graphviz（dot命令），您可以使用对应的Graphviz Python包（pip install graphviz）和一个简单的Python脚本来生成结果。

1. 我们将把 gprof2dot 的输出保存到一个 .dot 文件中： python -m gprof2dot -f pstats file.prof > file.dot
2. 我们将使用以下代码从这个.dot文件生成图片：

 import graphviz
     
 def make_png(input_file_name, output_file_name):
     dot = graphviz.Source.from_file(input_file_name)
     dot.render(outfile=output_file_name)
     
 if __name__ == '__main__':
     make_png('file.dot', 'file.svg')  # 还支持 .png、.jpg 等格式。

代码改进

让我们通过将两个正则表达式合并为一个来加快代码速度：

 def single_pattern():
     num_total_matches = 0
     pattern = r"[0-9][rb]ob"  # 我们现在只有一个搜索
     
     with open('random.txt', 'rt') as f:
         for line in tqdm(f, total=1_000_000):
             line = line.lower()
             # 在字符串中找到模式的所有匹配项
             matches = re.findall(pattern, line)
 ​
             # 计算比赛的数量
             num_matches = len(matches)
             num_total_matches += num_matches
         
     return num_total_matches

现在我们可以达到15秒，速度提高了约2倍！

新代码的性能分析结果。请注意findall函数所占用时间的百分比下降了（从88% -> 81.5%）

对新代码进行性能分析证实这是正确的选择 — 主要findall函数所占用时间比例下降，而其他所有函数所占用时间增加（大约增加了2倍）。

 def wasted_efforts():
     num_total_matches = 0
     # 我们使用编译后的正则表达式
     pattern1 = re.compile(r"[0-9]{1}rob", flags=re.IGNORECASE)
     # 我们还使用不区分大小写的情况。
     pattern2 = re.compile(r"[0-9]{1}bob", flags=re.IGNORECASE)
     
     with open('random.txt', 'rt') as f:
         # 我们将行加载到一个块中，并一次性处理。
         chunk = []
         for line in tqdm(f, total=1_000_000):
             chunk.append(line)
             
             if len(chunk) == 1000:
                 chunk_str = ''.join(chunk)
                 for pattern in [pattern1, pattern2]:
                     # 在字符串中找到模式的所有匹配项
                     matches = re.findall(pattern, chunk_str)
 ​
                     # 计算数量
                     num_matches = len(matches)
                     num_total_matches += num_matches
                     
                 chunk = []
 ​
     # 尽管我们修改了代码，但这段代码仍然和原始版本一样慢。
     return num_total_matches

多进程处理

findall 函数仍然是瓶颈。然而，在这个示例中，我们只能做一件直接的事情来进一步改进它：并行执行我们的代码。

不同字符串中的匹配是相互独立的，因此我们可以并行搜索它们。我们该如何做呢？

最简单的方法是创建一个进程池 — 一个操作多个并行 Python 进程的对象，所有进程都同时运行。

如果我们有一个值列表和要应用于每个值的函数，只需调用池的 map 方法，它将并行地在所有值上为我们运行函数。

 from multiprocessing import Pool
 ​
 def calc_num_matches(string):
     # 我们为我们的池对象创建一个单独的函数
     # 并行应用于文件中的每一行
     pattern = r"[0-9]{1}[rb]ob"
     return len(re.findall(pattern, string))
 ​
 def chunks_single_pool():
     num_total_matches = 0
     
     pool = Pool(8)  # 进程数量应该小于或等于 CPU 核心数。
     
     with open('random.txt', 'rt') as f:
         chunk = []
         # 我们将行加载到列表中
         for line in tqdm(f, total=1_000_000):
             line = line.lower()
             chunk.append(line)
             
             if len(chunk) == 1000:
                 # 然后我们并行地将我们的 `calc_num_matches` 函数应用于每一行
                 num_ind_matches = pool.map(calc_num_matches, chunk)
                 # 将所有独立匹配项相加。
                 num_matches = sum(num_ind_matches)
                 num_total_matches += num_matches
 ​
                 chunk = []
     
     return num_total_matches

让我们看看……6.1秒！这是一个新纪录。比我们优化的解决方案快2.2倍，比原始版本快约5倍！让我们来看一下性能分析结果：

多进程代码的性能分析结果。请注意正则表达式搜索所占比例不断减少（绿色块）。

Snakeviz显示多进程代码的相同图片。确认正则表达式搜索现在占据总时间的大约2/3，比原始版本少得多。

注意事项

多进程

需要注意是，分析器不再知道子进程中进行正则表达式搜索时发生了什么。它只显示等待所有子进程完成花费了4秒的时间，但由于它们基本上是与我们主要的Python进程松散连接的独立程序，因此无法访问在子进程中执行的代码。

如果你想分析程序中子进程的运行情况，应该将@profile装饰器放在将由子进程执行的函数/代码内部。

多线程

通常情况下，由于全局解释器锁（Global Interpreter Lock，GIL）的限制，不必直接编写多线程 Python 代码，因为 GIL 一次只允许一个线程执行 Python 代码（在单个 Python 进程内）。

然而，如果你在使用大量非Python库（例如NumPy、PyTorch、scipy）或进行大量I/O操作（网络通信），那么你的代码运行时大部分时间将花费在Python解释器之外执行用C、C++、Fortran等语言编写的代码上。

在这两种情况下，利用多线程可能是实际的，因此需要意识到 Python 的两个内置性能分析模块 — profile 和 cProfile — 仅对应用程序的主线程进行性能分析。如果想要对除主线程之外的其他线程执行的代码进行性能分析，可以尝试在将要执行线程的函数中运行分析器，或者使用一些第三方性能分析工具如 VizTracer。

总结

本文我们以示例演示一个存在性能问题的代码，并使用工具分析的场景。

在处理数据时遇到代码运行缓慢的问题，需要找出瓶颈所在并进行优化。通过生成大文本文件并定义基准函数来计算特定模式的出现次数，初始运行时间较长，需要加快速度。

解决方案：

1. 使用 SnakeViz 和 gprof2dot 等性能分析工具，找出主要瓶颈为 findall 函数中的正则表达式搜索。
2. 将两个正则表达式合并为一个，使运行速度提高约 2 倍。
3. 采用多进程处理，创建进程池并行执行代码，使运行速度比优化后的解决方案快 2.2 倍，比原始版本快约 5 倍。

注意事项：

1. 多进程方面，分析器无法知晓子进程中的正则表达式搜索情况，若要分析子进程，应将装饰器放在子进程执行的函数内。
2. 多线程方面，由于全局解释器锁限制，通常不必直接编写多线程 Python 代码。若使用大量非 Python 库或大量 I/O 操作，利用多线程可能有效。Python 内置性能分析模块仅对应用程序的主线程进行性能分析，若要分析其他线程，可在相关函数中运行分析器或使用第三方工具如 VizTracer。

作者：demo007x
链接：https://juejin.cn/post/7390319395833380901