神操作!两行代码,提速 13 倍!让 Python 飞一般的感觉!

开发 开发工具
Numba 是一款为 python 打造的、专门针对 Numpy 数组循环计算场景的即时编译器。显然,这正是我们所需要的。

Python 本身是一门运行较慢的语言,因此对于计算场景,最好的优化方式就是优化代码写法。你可以使用现有的科学计算库:比如 Numpy 和 Scipy。但如果想要在不使用低级语言(如 CPython、Rust 等)实现扩展的前提下实现一个新的算法时,该如何做呢?

对于某些特定的、尤其是针对数组的计算场景,Numba 可以显著加快代码的运行速度。在使用时,我们有时候需要调整一下原始代码,而有时候却又不需要做任何改动。当它真正起到作用时,效果将会非常明显。

在本篇文章中,我们会谈及以下几方面:

  • 为什么 有时候单独使用 Numpy 是不够的
  • Numba 的基础使用方式
  • Numba 是如何在很高的层次上来对你的代码运行造成影响的

Numpy ”爱莫能助“的时刻

假设你想要将一个非常大的数组转变为按递增顺序排序:很好理解,就是将元素按值的大小升序排列,如:

[1, 2, 1, 3, 3, 5, 4, 6]  [1, 2, 2, 3, 3, 5, 5, 6]

以下是一个简单的就地转换方式:

def monotonically_increasing(a):
max_value = 0
for i in range(len(a)):
if a[i] > max_value:
max_value = a[i]
a[i] = max_value

Numpy 运行很快,是因为它可以在不调用 python 自身解释器的前提下完成所有计算。但对于上面这个场景(python 中的循环),就会暴露出一个问题:我们会失去 Numpy 得天独厚的性能优势。

对一个含有一千万个元素的 Numpy 数组使用上面的函数进行转换,在我的电脑上需要运行 2.5 秒。那么,还可以优化得更快吗?

使用 Numba 提速

Numba 是一款为 python 打造的、专门针对 Numpy 数组循环计算场景的即时编译器。显然,这正是我们所需要的。让我们在原有函数的基础上添加两行代码试试:

from numba import njit

@njit
def monotonically_increasing(a):
max_value = 0
for i in range(len(a)):
if a[i] > max_value:
max_value = a[i]
a[i] = max_value

再次运行,发现仅需要 0.19 秒,在完全重用旧代码逻辑的前提下,感觉效果还不错。

实际上 Numpy 也有一个特殊的函数可以解决这种场景(但是会修改原有函数的代码逻辑):`numpy.maximum.accumulate` 。通过使用它,函数的运行时长会缩短至 0.03 秒。

Numba 简介

在 Numpy 或 Scipy 中找到目标函数,可以很快解决常见的计算问题。但是如果函数不存在呢?(比如刚刚的 numpy.maximum.accumulate)。这种情况下如果想加速代码运行。可能会选择其他低级的编程语言来实现扩展,但这也意味着切换编程语言,会让模块构建和系统总体变得更复杂。

使用 Numba 你可以做到:

  • 使用 python 和拥有更快编译速度的解释器运行同一份代码
  • 简单快速地迭代算法

Numba 首先会解析代码,然后根据数据的输入类型以即时的方式编译它们。例如,当输入是 u64 数组和浮点型数组时,分别得到的编译结果是不一样的。

Numba 还可以对非 CPU 的计算场景生效:比如你可以 在 GPU 上运行代码。诚然,上文中的示例只是 Numba 的一个最小应用,官方文档中还有很多特性可供选择。

Numba 的一些短板

需要一次代码编译耗时

当第一次调用 Numba 修饰的函数时,它需要花费一定的时间来生成对应的机器代码。比如,我们可以使用 IPython 的 %time 命令来计算运行一个 Numba 修饰的函数需要花费多长时间:

In [1]: from numba import njit

In [2]: @njit
...: def add(a, b): a + b

In [3]: %time add(1, 2)
CPU times: user 320 ms, sys: 117 ms, total: 437 ms
Wall time: 207 ms

In [4]: %time add(1, 2)
CPU times: user 17 µs, sys: 0 ns, total: 17 µs
Wall time: 24.3 µs

In [5]: %time add(1, 2)
CPU times: user 8 µs, sys: 2 µs, total: 10 µs
Wall time: 13.6 µs

可以看到,函数第一次调用后运行非常慢(注意单位时毫秒而不是微秒),这就是因为它需要时间来编译生成机器代码。不过函数后面的运行速度会显著提升。这种时间成本在输入数据的类型发生变化时会再次消耗,比如,我们将输入类型换为浮点数:

In [8]: %time add(1.5, 2.5)
CPU times: user 40.3 ms, sys: 1.14 ms, total: 41.5 ms
Wall time: 41 ms

In [9]: %time add(1.5, 2.5)
CPU times: user 16 µs, sys: 3 µs, total: 19 µs
Wall time: 26 µs

计算两数之和当然不需要启用 Numba,这里用这个案例是因为能够比较容易地看出编译所需的时间成本。

与 python 和 Numpy 的不同实现方式

Numba 在功能方面可以说是实现了 python 的一个子集,也可以说是实现了 Numpy API 的一个子集,这将会导致一些潜在的问题:

(1)会出现 python 和 Numpy 部分特性都不支持的情况

(2)由于 Numba 重新实现了 Numpy 的 API,在使用时可能会出现以下情况:

  • 由于使用的不用的算法,两者的性能表现会有区别
  • 可能会由于 bug 导致结果不一致

(3)另外,当 Numba 编译失败时,其暴露的错误信息可能会很难理解

Numba 与其他选项的对比

  • 仅使用 Numpy 和 Scipy:可以让 python 代码运行时达到其他语言编译器的速度,但是对于某些循环计算的场景不生效
  • 直接使用低级语言编写代码:这意味着你可以优化所有的代码语句,但是需要抛弃 python 使用另一门语言
  • 使用 Numba:可以优化 python 循环计算的场景,但是对于某些 python 语言本身和 Numpy API 的特性使用会受到限制

结语

Numba 最棒的地方在于尝试起来非常简单。因此每当你有一个做一些数学运算且运行缓慢的 for 循环时,可以尝试使用 Numba :运气好的话,它只需要两行代码就可以显著加快代码运行速度。

责任编辑:赵宁宁 来源: 菜鸟学Python
相关推荐

2024-03-28 18:05:41

Python编程

2020-04-26 17:04:18

Python代码数据

2022-03-18 07:44:59

插件VSCode代码

2018-07-27 09:32:18

Python代码数据

2020-07-20 09:20:48

代码geventPython

2021-06-18 10:12:09

JS代码前端

2018-09-19 15:46:51

编程语言Python编译器

2023-09-12 14:58:00

Redis

2013-02-28 10:35:59

hadoop大数据Hortonworks

2016-03-21 10:16:06

RedisSpark大数据处理

2022-09-25 23:10:53

Python数据集机器学习

2019-09-09 15:18:30

物联网5G技术

2020-07-22 08:30:02

代码开发工具

2023-12-11 15:40:32

PyTorch代码大模型

2016-10-08 16:02:37

WIFIMegaMIMO系统

2023-03-16 16:18:09

PyTorch程序人工智能

2015-07-22 15:53:57

2024-02-20 12:49:00

CSS函数前端

2017-08-31 13:50:53

Python编程语言

2023-08-14 07:42:01

模型训练
点赞
收藏

51CTO技术栈公众号