目录
分析你的 Numba 代码
如果你正在编写数值计算的 Python 代码，Numba 可以极大地加速你的程序。通过将 Python 的一个子集编译为机器码，Numba 允许你编写在普通 Python 中运行过慢的 for 循环和其他结构。换句话说，它类似于 Cython、C 或 Rust，允许你为 Python 编写编译扩展。
然而，Numba 代码并不总是能达到最快的速度。这时，性能分析就派上用场了：它可以帮助你找到代码中的瓶颈。
介绍 Profila：Numba 的性能分析工具
Profila 是一个专门用于识别 Numba 代码中哪些行运行缓慢的性能分析工具。它通过采样工作：它运行一个 gdb 进程，每隔 10 毫秒左右连接到你的进程，获取回溯信息，并使用这些信息来识别哪些代码行正在运行。获取足够的样本后，你就能了解程序中哪些部分花费了最多的时间。
让我们看看 Profila 的实际应用！我们将分析以下程序，该程序对图像进行抖动处理，将其从灰度图像转换为黑白图像。
from numba import njit
import numpy as np
from skimage import io
from time import time
IMAGE = io.imread("hallway.jpg")
# 使用 numba.njit 装饰的代码看起来像 Python 代码，
# 但实际上是在运行时编译为机器码的快速低级代码！
@njit
def dither(img):
# 允许负值和比 uint8 更宽的范围：
result = img.astype(np.int16)
y_size = img.shape[0]
x_size = img.shape[1]
last_y = y_size - 1
last_x = x_size - 1
for y in range(y_size):
for x in range(x_size):
old_value = result[y, x]
if old_value < 0:
new_value = 0
elif old_value > 255:
new_value = 255
else:
new_value = np.uint8(np.round(old_value / 255.0)) * 255
result[y, x] = new_value
# 我们可能会得到一个负值的误差：
error = np.int16(old_value) - new_value
if x < last_x:
result[y, x + 1] += error * 7 // 16
if y < last_y and x > 0:
result[y + 1, x - 1] += error * 3 // 16
if y < last_y:
result[y + 1, x] += error * 5 // 16
if y < last_y and x < last_x:
result[y + 1, x + 1] += error // 16
return result.astype(np.uint8)
# 第一次运行以编译代码：
dither(IMAGE)
# 确保我们运行足够长的时间，以便性能分析器获取足够的样本：
start = time()
runs = 0
while time() - start < 5:
runs += 1
dither(IMAGE)
elapsed = time() - start
print(f"Processed {int(round(runs / elapsed))} images / sec")
然后我们安装 gdb 和 profila：
$ sudo apt-get install -y gdb
$ pip install profila
现在我们可以使用 Profila 来分析代码：
$ python -m profila annotate -- dither.py
Processed 255 images / sec
# Total samples: 805 (37.8% non-Numba samples, 1.9% bad samples)
## File `dither.py`
Lines 14 to 40:
0.1% | result = img.astype(np.int16)
| y_size = img.shape[0]
| x_size = img.shape[1]
| last_y = y_size - 1
| last_x = x_size - 1
| for y in range(y_size):
0.2% | for x in range(x_size):
| old_value = result[y, x]
| if old_value < 0:
| new_value = 0
1.0% | elif old_value > 255:
| new_value = 255
| else:
30.2% | new_value = np.uint8(np.round(old_value / 255.0)) * 255
| result[y, x] = new_value
| # 我们可能会得到一个负值的误差：
2.0% | error = np.int16(old_value) - new_value
0.6% | if x < last_x:
11.3% | result[y, x + 1] += error * 7 // 16
| if y < last_y and x > 0:
0.9% | result[y + 1, x - 1] += error * 3 // 16
| if y < last_y:
2.4% | result[y + 1, x] += error * 5 // 16
| if y < last_y and x < last_x:
1.6% | result[y + 1, x + 1] += error // 16
|
0.2% | return result.astype(np.uint8)
## File `numba/np/arraymath.py`
Lines 3157 to 3157:
9.8% | return _np_round_float(a)
优化代码
首先，从输出的顶部可以看到，37.8 + 1.9 = 大约 40% 的样本要么是无效的，要么没有运行 Numba 代码。这包括导入模块、运行常规 Python 代码、编译 Numba 代码等时间。这意味着 60% 的时间用于运行实际的 Numba 代码。
然后，我们可以看到 np.round() 和 new_value = np.uint8(np.round(old_value / 255.0)) * 255 的组合占用了 30.2 + 9.8 = 40% 的总样本。这意味着 40/60，即三分之二的时间，花费在这一行代码上，将 old_value 转换为 new_value。
可以将以下代码：
if old_value < 0:
new_value = 0
elif old_value > 255:
new_value = 255
else:
new_value = np.uint8(np.round(old_value / 255.0)) * 255
替换为一个更简单的版本：
new_value = 0 if old_value < 128 else 255
如果我们分析这个优化后的版本，我们会看到：
$ python -m profila annotate -- dither2.py
Processed 758 images / sec
# Total samples: 713 (30.3% non-Numba samples, 2.2% bad samples)
## File `/home/itamarst/devel/sandbox/numba-profiling-article/dither2.py`
Lines 14 to 35:
0.1% | result = img.astype(np.int16)
| y_size = img.shape[0]
| x_size = img.shape[1]
| last_y = y_size - 1
| last_x = x_size - 1
| for y in range(y_size):
0.3% | for x in range(x_size):
0.1% | old_value = result[y, x]
4.5% | new_value = 0 if old_value < 128 else 255
2.0% | result[y, x] = new_value
| # 我们可能会得到一个负值的误差：
5.5% | error = np.int16(old_value) - new_value
5.3% | if x < last_x:
34.9% | result[y, x + 1] += error * 7 // 16
| if y < last_y and x > 0:
3.6% | result[y + 1, x - 1] += error * 3 // 16
| if y < last_y:
6.0% | result[y + 1, x] += error * 5 // 16
| if y < last_y and x < last_x:
4.5% | result[y + 1, x + 1] += error // 16
|
0.6% | return result.astype(np.uint8)
我们的代码现在运行速度快了 3 倍！
这一次，我们不需要猜测，而是有了一个性能分析工具来指出主要的瓶颈。
性能分析的局限性
局限性 1：编译后的代码并不总是与源代码匹配
考虑以下示例：
@njit
def add(arr):
result = np.empty_like(arr)
for i in range(len(arr)):
orig = arr[i]
orig += 123_000_000
orig += 456_000
orig += 789
result[i] = orig
return result
如果我们分析它，我们会得到：
## File `combo.py`
Lines 13 to 14:
2.9% | orig += 789
7.5% | result[i] = orig
为什么只列出了 orig += 789，而没有列出其他添加 123_000_000 和 456_000 的行？这是因为编译器优化，它们在编译版本中不再作为独立的操作存在。我们可以检查生成的汇编代码并搜索 789：
>>> for line in next(iter(add.inspect_asm().values())).splitlines():
... if "789" in line: print(line)
...
.long 123456789
movl $123456789, %edi
movl $123456789, %edx
编译器将三个相加的数字（123_000_000、456_000、789）合并为一个数字（123_456_789），这可能是通过一次加法操作完成的，而不是三次。
局限性 2：你需要对 CPU 执行方式有一个心理模型
回想一下，在我们优化的抖动函数中，我们看到了以下性能分析结果：
5.3% | if x < last_x:
34.9% | result[y, x + 1] += error * 7 // 16
| if y < last_y and x > 0:
3.6% | result[y + 1, x - 1] += error * 3 // 16
| if y < last_y:
6.0% | result[y + 1, x] += error * 5 // 16
| if y < last_y and x < last_x:
4.5% | result[y + 1, x + 1] += error // 16
为什么第一个计算比后面的计算要昂贵得多？它们似乎在做几乎相同的事情。
一个假设是指令级并行性：CPU 能够提前并行执行后面的代码。因此，第一行之所以昂贵，并不是因为它有什么不同，而是因为它是第一个。我们可以通过改变语句的顺序并再次分析来部分验证这个假设。以下是结果：
| if y < last_y and x > 0:
20.3% | result[y + 1, x - 1] += error * 3 // 16
14.2% | if x < last_x:
8.8% | result[y, x + 1] += error * 7 // 16
| if y < last_y:
3.7% | result[y + 1, x] += error * 5 // 16
| if y < last_y and x < last_x:
3.7% | result[y + 1, x + 1] += error // 16
现在，result[y, x + 1] += error * 7 // 16 这一行的样本百分比比之前的代码版本低得多，即使它运行的次数和计算完全相同。唯一改变的是执行顺序。
一个仅线性执行代码、没有并行性的 CPU 心理模型不足以理解这些结果。
局限性 3：更快的代码可能需要更全面的视角
看到哪些代码行运行缓慢非常有帮助，但一些优化是重新思考代码结构的结果。例如，最终优化版本将内存使用量减少了三分之二，并且比第二个版本快了 50%。它通过改变算法的结构来实现这一点，而不是仅仅关注一两行代码。
扫一扫
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
