Python矢量化作为一种性能技术的局限性

吴欣
403 阅读10分钟

正如NumPy所实现的那样，Python中的矢量化可以通过使用快速的低级代码对大量数据进行操作，给你带来更快的操作。而Pandas建立在NumPy的基础上，提供类似的快速功能。但是，矢量化并不是解决你所有问题的灵丹妙药：有时它将以更高的内存使用率为代价，有时你需要的操作不被支持，有时它只是不相关。

对于每一个问题，都有其他的解决方案可以解决这个问题。在这篇文章中，我们将

  1. 回顾一下为什么矢量化是有用的。
  2. 复述矢量化的各种限制和问题。
  3. 考虑每个问题的一些解决方案：PyPy、Numba和编译的扩展。

矢量化有多种含义，我们的重点是快速、低级的循环

总结一下对矢量化更详细的解释，在Python的上下文中，这个词有三种不同的含义。

  1. 一个对批量数据进行操作的API。例如，arr += 1 会给NumPy数组中的每一项添加1。
  2. 一个用低级语言（C、Rust）实现的快速API，可以对大量数据进行快速操作。这将是我们在本文中的主要重点。
  3. 利用SIMD CPU指令来进一步加快低级别的操作。详见矢量化的主要写法。

我们将专注于第二个意思，从假设默认的Python解释器开始，称为CPython。当你在命令行上运行python ，你通常在运行CPython。

一个例子：将一个整数添加到一个整数列表中

考虑一下下面的例子。

from time import time

l = list(range(100_000_000))

start = time()
for i in range(len(l)):
    l[i] += 17
print("Elapsed: ", time() - start)

一个由 Python 整数组成的 CPython 列表是一系列指向相当大的对象的指针，从 CPython 的角度来看，这些对象只是一般的 Python 对象。为 Python 列表中的每一个项目添加一个整数涉及到。

  1. 查找列表中每个项目的类型。
  2. 查找该类型的 add 函数。
  3. 用原始对象和被添加的 Python 对象调用该函数。
  4. 将两个 Python 对象转换成机器码的整数。
  5. 添加机器码的整数。
  6. 将得到的机器码整数包入一个 Python 对象。
  7. 在 Python 列表中查找下一个项目，这本身就是一个昂贵的操作，涉及到查找范围类型的迭代器方法，然后查找列表类型的索引操作，然后将索引对象转换为机器码的整数...

这是一个很大的工作!

相比之下，NumPy的整数数组就是......一个机器编码的整数数组。因此，给每个项目添加一个整数只需要少量的CPU指令，一旦初始设置代码完成，就和同等的C代码没有区别。

下面是NumPy的代码。

from time import time
import numpy as np

l = np.array(range(100_000_000), dtype=np.uint64)

start = time()
l += 17
print("Elapsed: ", time() - start)

正如你所期望的，NumPy更快。快得多。

执行耗时(秒)
CPython6.13
NumPy0.07

矢量化的局限性

矢量化使你的代码更快，这很好......但它不是一个完美的解决方案。让我们考虑一下其中的一些问题。

问题#1：不必要的大内存分配

让我们看看你想计算一个数组的项与0的平均距离。一种方法是计算每个项的绝对值，然后计算结果的平均值。

imoprt numpy as np

def mean_distance_from_zero(arr):
    return np.abs(arr).mean()

这很有效，而且这两种计算方法都很快速，而且是矢量的......但是这涉及到创建一个全新的中间数组来存储绝对值。如果你的原始数组是1000MB，你刚刚又分配了1000MB。你可以覆盖初始数组以节省内存，但你可能因为其他原因需要它。

在实践中，你可以用原地操作来限制它，这样就只有一个额外的副本，而不是N个副本。还有一些库，如 numexpr和Dask Array这样的库，可以为你做批处理和更智能的原地更新，以减少内存负担。

另外，你也可以用自己的一个循环来逐项计算......但这就导致了我们的第二个问题。

问题二：只有支持的操作才是快速的

为了使矢量化工作，你需要低级别的机器代码来驱动数据的循环，并运行实际的操作。换回Python的循环和功能就会失去这种速度。例如，对于我们上面的mean_distance_from_zero() 函数，在不占用较大内存的情况下，实现它的明显方法是这样的。

def mean_distance_from_zero(arr):
    total = 0
    for i in range(len(arr))
        total += abs(arr[i])
    return total / len(arr)

但是我们又回到了在Python中做循环，在Python中做操作；代码又很慢了。这就是为什么NumPy必须自己实现这么多东西的原因之一：任何时候你必须回到Python，事情就会变得很慢。作为一个额外的惩罚，在一个编译语言中，编译器也有可能找到一些方法来优化我们的第二个实现，给我们一个加速，而这是一系列单独的矢量化操作所不能做到的。

问题三：矢量化只能加快批量操作的速度

有时你的代码很慢，因为它对许多相同类型的数据项做了完全相同的操作。在这种情况下，矢量化是你的朋友。

在其他情况下，你的计算太慢了，但并不符合那个特定的模式。在这一点上，矢量化并不能帮助你，你需要一些其他的解决方案来加速你的代码。

一些解决方案

这里我们将介绍三种基本的潜在解决方案。

  1. 一个更快的 Python 解释器。PyPy是CPython的另一种实现，它更聪明：它可以使用即时编译，通过生成适当的、专门的机器代码来加快执行速度。
  2. 通过使用Numba，在CPython本身中生成机器代码。
  3. 使用Cython、C、C++或Rust编译自定义代码。

PyPy：一个更快的Python解释器

如果我们有一个Python整数的列表，CPython（默认的Python解释器）并没有真正对它做任何聪明的处理。这就是为什么我们需要使用NumPy数组。PyPy更聪明；它可以为这种情况生成专门的代码。再回到我们最初的例子。

from time import time

l = list(range(100_000_000))

start = time()
for i in range(len(l)):
    l[i] += 17
print("Elapsed: ", time() - start)

我们可以比较CPython和PyPy。

$ python add_list.py
Elapsed:  6.125197410583496
$ pypy add_list.py
Elapsed:  0.11461925506591797

PyPy几乎和NumPy的代码一样快，而且......它只是一个普通的Python循环。因为它知道你只有一个整数的列表，所以它可以为这种特殊情况生成专门的、快得多的机器代码。

不幸的是，PyPy与NumPy的接口不是特别好；在NumPy数组上的一个天真的非矢量循环（比如我们上面实现的mean_distance_from_zero() ），在PyPy中实际上比CPython慢得多。因此，如果你想处理NumPy中缺失的功能，PyPy是没有帮助的；它甚至使缓慢的无矢量循环解决方案变得更慢。

然而，如果你在处理使用标准Python对象的代码，PyPy可以快得多。所以它最适合用于加速我们完全不使用NumPy或Pandas的情况。

Numba：与NumPy一起使用的及时函数

Numba也可以进行及时编译，但与PyPy不同的是，它是作为标准CPython解释器的一个插件--它是为NumPy设计的。让我们看看它是如何解决我们的问题的。

用Numba扩展NumPy

使用Numba，缺失的操作不是问题；你可以直接写你自己的。例如，让我们用纯Python和Numba来实现我们的mean_distance_from_zero() 。

from time import time
import numpy as np
from numba import njit

# Pure Python version:
def mean_distance_from_zero(arr):
    total = 0
    for i in range(len(arr))
        total += abs(arr[i])
    return total / len(arr)

# A fast, JITed version:
mean_distance_from_zero_numba = njit(
    mean_distance_from_zero
)

arr = np.array(range(10_000_000), dtype=np.float64)

start = time()
mean_distance_from_zero(arr)
print("Elapsed CPython: ", time() - start)

for i in range(3):
    start = time()
    mean_distance_from_zero_numba(arr)
    print("Elapsed Numba:   ", time() - start)

结果。

$ python cpython_vs_numba.py
Elapsed CPython:  1.1473402976989746
Elapsed Numba:    0.1538538932800293
Elapsed Numba:    0.0057942867279052734
Elapsed Numba:    0.005782604217529297

像PyPy一样，Numba为这个函数生成了专门的机器码，尽管与PyPy不同，它只能为Python语言的一个子集这样做。对Numba的第一次调用要慢得多，因为Numba要编译一些自定义的机器代码；之后的调用就特别快。而正如预期的那样，CPython要比Numba慢得多。

因为我们可以使用Numba为NumPy编写自定义的附加函数，所以缺失的操作并不是问题。这也意味着我们不会因为NumPy可用数组的限制而被迫创建不必要的临时数组，所以Numba可以帮助我们解决我们看到的前两个问题。

Numba对非矢量操作的帮助较小，因为它们不是项目目标的一部分。

用Cython（或Rust，或C，或C++）编译的代码

到目前为止，我们已经看到了两个及时编译的例子：整个解释器，或者只是函数。另一种加快代码的方法是提前编译代码。使用Cython，或Rust，或C，你可以。

  1. 使用Cython的原生NumPy支持、Rust的NumPy支持为NumPy编写快速、额外的操作。 rust-numpy，或者直接使用Python C API。NumPy本身大部分是用C语言编写的，现有的NumPy扩展也是用其他语言编写的，如Fortran或C++。
  2. 对于非矢量的用例，你可以再次使用这些语言为Python编写快速扩展。

然而，你必须在你的打包或构建设置中添加一堆配置，以便在Python运行前编译这些扩展。Cython实际上有一个在导入时编译的选项，但这使得你的软件更难发布，因为它要求用户安装一个编译器。Numba则没有这种限制。

选择一个解决方案

对于矢量操作的用例，PyPy是不会有帮助的。像HPy项目这样正在进行的努力可能会使它在未来更加有用。

这就剩下Numba和编译扩展了。

  • 因为你需要提前编译，这需要一个编译器和更复杂的包装，Numba通常更容易设置。
  • Numba可以为不同类型的整数或浮点数建立不同版本的代码，而不需要额外的工作；用编译的代码，你必须为每个版本编译一个版本，也许还要为每个版本写代码。对于Rust或C++，你可以通过泛型或模板来解决代码重复的问题。
  • Numba是Python的一个非常有限的子集，调试它的局限性会很困难。相比之下，编译语言要灵活得多；Cython几乎支持所有的Python，Rust是一种复杂的语言，有优秀的工具，等等。

对于非矢量操作，PyPy通常会让你的代码更快，而不需要做任何改变；你可以直接尝试你的代码，看看是否更快。否则，你就只能对现有的代码进行优化，或者缺乏选择合适的编译语言来写一个编译扩展。

avatar
文章
阅读
粉丝