Python代码性能优化的神器：Numba
Python代码性能优化，Numba绝对是个神器！它能把你的Python代码直接编译成机器码，速度能提升好几倍，有时候甚至能提升几百倍。最爽的是，你都不用改多少代码，加个装饰器就搞定了。它特别适合处理数值计算、科学计算这些耗时的操作，就像给代码装了个超级加速器。
安装和基本使用
装Numba贼简单，一行命令搞定：
pip install numba
来看看最基础的用法，加个@jit装饰器就能提速：
from numba import jit
import numpy as np
@jit
def calculate_sum(arr):
total = 0
for i in range(len(arr)):
total += arr[i]
return total
# 测试一下速度
data = np.arange(10000000)
%time result = calculate_sum(data) # 要在Jupyter里运行才能看到时间
温馨提示：第一次运行可能会慢一点，这是因为Numba在编译代码，后面再运行就飞快了。
并行计算加速
想更快？来试试并行计算：
from numba import prange, jit
@jit(parallel=True)
def parallel_sum(arr):
total = 0
for i in prange(len(arr)): # 用prange替换range
total += arr[i]
return total
这代码看着差不多，但用了prange后，Numba就会自动把任务分给多个CPU核心去处理，速度刷刷地往上涨。
向量化运算
数组运算能玩出花来：
from numba import vectorize
import numpy as np
@vectorize
def fast_add(a, b):
return a + b
# 测试一下
x = np.arange(1000000)
y = np.arange(1000000)
result = fast_add(x, y)
温馨提示：vectorize装饰器特别适合处理大量数据的数组运算，它能自动把运算并行化。
GPU加速
要是有显卡，还能用GPU加速，那速度就更快了：
from numba import cuda
@cuda.jit
def gpu_add(arr_a, arr_b, result):
idx = cuda.grid(1)
if idx < result.size:
result[idx] = arr_a[idx] + arr_b[idx]
# 准备数据
a = np.arange(1000000)
b = np.arange(1000000)
result = np.zeros_like(a)
# 转移到GPU
d_a = cuda.to_device(a)
d_b = cuda.to_device(b)
d_result = cuda.to_device(result)
# 运行
threads_per_block = 256
blocks_per_grid = (a.size + threads_per_block - 1) // threads_per_block
gpu_add[blocks_per_grid, threads_per_block](d_a, d_b, d_result)优化注意事项
用Numba也有些坑得注意：
-
不是所有Python代码都能加速，主要是数值计算相关的代码效果最好 -
代码里尽量用NumPy数组，少用Python的列表 -
循环里面别整太多花里胡哨的操作 -
函数最好是纯函数，别用全局变量
看代码：
# 这样写效果好
@jit
def good_function(x):
return x * x + 2 * x + 1
# 这样写效果差
@jit
def bad_function(x):
result = []
for i in x:
result.append(str(i)) # 字符串操作不适合用Numba
return result
温馨提示：用之前先测试一下性能提升多少，别一股脑就上，有时候简单代码用Numba反而更慢。
代码性能优化就像给汽车改装，Numba就是个超级改装工具。不过也别啥代码都想着优化，把主要精力放在真正耗时的地方。写代码时把性能热点找准了，加个Numba装饰器，性能立马嗖嗖往上涨，这感觉不要太爽！
来源：
THE END
