Python高级技术之：如何利用`Numba`库，加速`Python`函数的执行。 - 智猿学院-前后端，数据库，人工智能，云计算等领域前沿技术讲座

观众朋友们，大家好！今天咱们来聊聊Python提速的秘密武器之一：Numba。别害怕，虽然听起来像什么魔法咒语，但其实它很简单，就像给你的Python代码喝了红牛，瞬间充满能量！

一、Numba：你的Python代码加速器

Python很棒，但速度嘛…有时候像蜗牛散步。特别是当你的代码涉及到大量的循环和数学运算时，它可能会让你等到天荒地老。这时候，Numba就派上用场了。

Numba是一个开源的JIT（Just-In-Time）编译器，它可以将你的Python函数“编译”成机器码，从而大大提高运行速度。注意，这里说的是“编译”，但不是像C++那样提前编译好，而是在运行时，根据你的代码和数据类型，动态地生成机器码。

二、Numba的原理：JIT编译

JIT编译就像一个翻译官，它不是提前把所有东西都翻译好，而是当你需要的时候，才把相关的部分翻译成机器能听懂的“语言”。这样既灵活，又高效。

具体来说，Numba会分析你的Python函数，找出可以加速的部分，然后将这些部分编译成机器码。这个过程是在运行时发生的，所以它能够根据实际的数据类型进行优化。

三、安装Numba：很简单，一键搞定

想要使用Numba，首先要安装它。很简单，打开你的终端或命令提示符，输入以下命令：

pip install numba

搞定！是不是感觉像变魔术一样？

四、Numba的基本用法：加个装饰器就行

使用Numba最简单的方法就是给你的函数加上一个装饰器@jit。就像给你的函数贴了一张“加速符”一样。

from numba import jit

@jit
def add_numbers(x, y):
  """
  一个简单的加法函数。
  """
  result = x + y
  return result

# 现在，当你调用 add_numbers 时，Numba 会尝试将其编译成机器码。
print(add_numbers(10, 20)) #输出30

这段代码中，@jit就是装饰器。它告诉Numba：“嘿，这个函数很重要，帮我把它加速一下！”。

五、Numba的两种模式：nopython 和 objectmode

Numba有两种主要的编译模式：nopython模式和objectmode模式。

nopython模式： 这是Numba最强大的模式。在这种模式下，Numba会尝试将整个函数都编译成机器码，不依赖于Python的解释器。如果编译成功，速度可以提升非常多。但如果Numba无法编译整个函数（比如函数中使用了Numba不支持的Python特性），就会报错。

要启用nopython模式，你可以使用@jit(nopython=True) 或者 @njit (它是 @jit(nopython=True) 的简写):
```
from numba import njit

@njit
def calculate_sum(arr):
    """
    计算数组元素的总和。
    """
    total = 0
    for i in range(arr.shape[0]):
        total += arr[i]
    return total

import numpy as np

my_array = np.arange(1000)
print(calculate_sum(my_array)) #输出499500
```
objectmode模式： 如果nopython模式失败了，Numba会自动退回到objectmode模式。在这种模式下，Numba会尽可能地编译函数，但仍然会依赖于Python的解释器。速度提升不如nopython模式那么明显，但至少比纯Python代码要快一些。

你也可以显式地指定objectmode模式：@jit(nopython=False)。
```
from numba import jit

@jit(nopython=False)
def process_data(data):
    """
    处理数据，包含一些Python对象操作。
    """
    results = []
    for item in data:
        results.append(item * 2) # 这里假设 data 包含的是可以进行乘法操作的对象
    return results

data = [1, 2, 3, 4, 5]
print(process_data(data)) #输出[2, 4, 6, 8, 10]
```
一般来说，我们应该尽量让Numba工作在nopython模式下，因为这样才能获得最大的速度提升。

六、Numba的优势和局限性

优势：

加速Python代码： 这是Numba最主要的优势。它可以显著提高你的Python代码的运行速度，特别是对于那些涉及到大量循环和数学运算的代码。
易于使用： 只需要加一个装饰器，就可以让Numba开始工作。
与NumPy集成： Numba与NumPy配合得非常好，可以加速NumPy数组的运算。
支持多种数据类型： Numba支持多种数据类型，包括整数、浮点数、复数等。
并行计算： Numba可以利用多核CPU进行并行计算，进一步提高速度。

局限性：

并非所有Python代码都能加速： Numba主要针对的是那些涉及到大量循环和数学运算的代码。对于那些主要涉及到I/O操作或者字符串处理的代码，Numba可能无法带来明显的加速效果。
对Python特性的支持有限： Numba并非支持所有的Python特性。如果你的代码中使用了Numba不支持的特性，Numba可能无法编译你的函数。
编译时间： Numba需要在运行时编译你的函数，这会带来一定的编译时间。不过，一旦函数被编译过，下次再调用时就可以直接使用编译后的代码，而不需要重新编译。
调试困难： 当Numba编译的代码出现问题时，调试起来可能会比较困难。

七、Numba实战：加速NumPy数组运算

NumPy是Python中用于科学计算的重要库。Numba与NumPy配合得非常好，可以加速NumPy数组的运算。

import numpy as np
from numba import njit
import time

@njit
def calculate_distance(x1, y1, x2, y2):
    """
    计算两点之间的距离。
    """
    return np.sqrt((x1 - x2)**2 + (y1 - y2)**2)

@njit
def calculate_all_distances(x, y):
    """
    计算所有点之间的距离。
    """
    n = len(x)
    distances = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            distances[i, j] = calculate_distance(x[i], y[i], x[j], y[j])
    return distances

# 生成随机坐标
n = 1000
x = np.random.rand(n)
y = np.random.rand(n)

# 计时：使用Numba加速的版本
start_time = time.time()
distances_numba = calculate_all_distances(x, y)
end_time = time.time()
numba_time = end_time - start_time
print(f"Numba加速版本耗时: {numba_time:.4f} 秒")

# 不使用Numba的版本 (为了公平比较，需要稍微修改代码，避免numba在第一次运行的时候编译)
def calculate_distance_no_numba(x1, y1, x2, y2):
    return np.sqrt((x1 - x2)**2 + (y1 - y2)**2)

def calculate_all_distances_no_numba(x, y):
    n = len(x)
    distances = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            distances[i, j] = calculate_distance_no_numba(x[i], y[i], x[j], y[j])
    return distances

start_time = time.time()
distances_no_numba = calculate_all_distances_no_numba(x, y)
end_time = time.time()
no_numba_time = end_time - start_time
print(f"不使用Numba版本耗时: {no_numba_time:.4f} 秒")

print(f"加速比: {no_numba_time / numba_time:.2f} 倍")

在这个例子中，我们定义了两个函数：calculate_distance用于计算两点之间的距离，calculate_all_distances用于计算所有点之间的距离。我们使用@njit装饰器来加速这两个函数。

运行这段代码，你会发现Numba加速后的版本比不使用Numba的版本快得多。加速比通常能达到几倍甚至几十倍。

八、Numba高级技巧：显式指定类型

有时候，Numba无法自动推断出你的变量类型。这时候，你可以显式地指定变量类型，以帮助Numba更好地进行编译。

from numba import njit, float64, int32

@njit(float64(int32, int32)) # 指定输入参数和返回值的类型
def multiply_numbers(x, y):
    """
    一个简单的乘法函数，显式指定类型。
    """
    return x * y

print(multiply_numbers(5, 10)) #输出50

在这个例子中，我们使用float64(int32, int32)来指定输入参数x和y的类型为int32，返回值的类型为float64。

常用的Numba类型包括：

类型	描述
`int32`	32位整数
`int64`	64位整数
`float32`	32位浮点数
`float64`	64位浮点数
`boolean`	布尔值
`complex64`	64位复数 (实部和虚部都是32位浮点数)
`complex128`	128位复数(实部和虚部都是64位浮点数)

九、Numba与并行计算：多核加速

Numba可以利用多核CPU进行并行计算，进一步提高速度。要启用并行计算，你需要设置parallel=True：

from numba import njit, prange
import numpy as np

@njit(parallel=True)
def calculate_sum_parallel(arr):
    """
    使用并行计算计算数组元素的总和。
    """
    n = len(arr)
    total = 0.0
    for i in prange(n): # 使用 prange 代替 range
        total += arr[i]
    return total

my_array = np.arange(100000)
print(calculate_sum_parallel(my_array)) #输出4999950000.0

在这个例子中，我们使用prange代替了range。prange是Numba提供的并行循环，它可以将循环任务分配给多个CPU核心并行执行。

需要注意的是，并行计算并非总是能带来速度提升。如果你的任务本身比较简单，或者CPU核心数量较少，并行计算可能会因为线程切换的开销而导致速度下降。

十、Numba的注意事项

预热： Numba的JIT编译需要一定的时间。第一次运行被@jit装饰的函数时，会比未被装饰的函数慢。但是后续的运行会非常快。因此，通常建议先“预热”一下，即先运行一次函数，让Numba完成编译，然后再进行实际的计算。
避免在Numba函数中使用Python对象： 尽量在Numba函数中使用NumPy数组和Numba支持的数据类型。避免使用Python的list、dict等对象，因为这些对象在Numba中处理起来比较慢。
小心副作用： Numba编译后的代码与Python代码的执行顺序可能不同。因此，要小心副作用，避免在Numba函数中修改全局变量或者进行I/O操作。

总结：

Numba是一个强大的Python加速工具，它可以显著提高你的Python代码的运行速度。掌握Numba的基本用法，可以让你在处理大规模数据和复杂计算时更加得心应手。但是，Numba并非万能的，你需要了解它的优势和局限性，才能更好地利用它。

希望今天的讲座对大家有所帮助！记住，写代码就像烹饪，Numba就像是你的秘密调料，用好了，能让你的代码美味无比！感谢大家的观看！

发表回复 取消回复

发表回复取消回复