请求提供使用Numba.jit()提升Python性能的示例

Numba的jit()装饰器核心作用是把Python函数编译成机器码，大幅提升计算密集型任务的运行速度，尤其适合循环多、数值运算多的场景。

基础用法示例 #

普通Python函数（无优化） #

先写一个计算从0到n-1的平方和的普通函数：

def sum_of_squares(n):
    total = 0
    for i in range(n):
        total += i ** 2
    return total

用jit()装饰优化后的函数 #

from numba import jit

# nopython=True是推荐模式，强制完全编译为机器码，不依赖Python解释器
@jit(nopython=True)
def numba_sum_of_squares(n):
    total = 0
    for i in range(n):
        total += i ** 2
    return total

性能对比测试 #

可以通过简单的计时代码直观看到优化效果：

import time

# 测试普通函数
start = time.time()
sum_of_squares(10_000_000)
print(f"普通函数耗时: {time.time() - start:.4f} 秒")

# 测试Numba函数：第一次调用会触发编译，耗时略长
start = time.time()
numba_sum_of_squares(10_000_000)
print(f"Numba函数首次调用耗时: {time.time() - start:.4f} 秒")

# 第二次调用直接用编译好的机器码，速度显著提升
start = time.time()
numba_sum_of_squares(10_000_000)
print(f"Numba函数二次调用耗时: {time.time() - start:.4f} 秒")

关键注意事项 #

• nopython=True的必要性：这个参数强制Numba生成纯机器码，避免回退到Python解释器，是性能提升的核心。如果不设置，Numba会尝试编译，失败则切换到对象模式，优化效果有限。
• 适用场景：只对计算密集型任务有效（比如数值模拟、矩阵运算、大量循环），IO密集型任务（比如文件读写、网络请求）几乎没有提升。
• 静态类型要求：nopython模式下，尽量避免Python的动态特性（比如动态修改列表长度、使用非预定义类型的字典键），尽量使用固定类型的变量和操作。
• 编译缓存：可以添加cache=True参数（@jit(nopython=True, cache=True)），把编译结果缓存到本地，下次运行脚本时无需重新编译。

内容的提问来源于stack exchange，提问作者Асмита Язь-взять