几种流行的 Python 性能加速方案对比

l1xnan
2,299 阅读4分钟

众所周知，Python 在几种常用的脚本语言（Ruby、JavaScript、Lua 等）中，速度是垫底的的存在。业界在上面也投入了很大精力，提供了各种各样的解决方案，各有各的优缺点，下面对这些解决方案做一些简单的介绍。

官方方案

首先不得不提的是Python之父参与的“香农计划”，在 2021 年 5 月的 Python 语言峰会上，Python 之父 Guido van Rossum 作了一场《Making CPython Faster 》的分享，宣布投入该计划，要在未来四年内将 CPython 速度提升 5 倍。 方案上涉及 CPython 一系列细节上的改进，有兴趣的可以自行搜索，最新进展是在 Python 3.11 版本，提速 10-60%。 作为用户只需等官方进展，升级 Python 版本既可。

PyPy

PyPy 是一种快速、兼容的 Python 替代实现，得益于其 Just-in-Time 编译器，Python 程序在 PyPy 上通常运行得更快。官方提供的基准测试相比 CPython 3.7 快 4.8 倍。安装使用上和官方的 Python 一样。最新版本 PyPy 3.9，相比官方的 Python 3.11 新特性可能不支持。

Cython

Cython 是一种针对 Python 和扩展的 Cython(基于 Pyrex)的优化静态编译器。它可以让我们为 Python 编写 C 扩展时，像编写 Python 本身一样简单。最著名的项目就是 numpy。 与 Python 不同，Cython 代码必须被编译。这发生在两个阶段：

  • Cython 将 .pyx 或 .py 文件编译为 .c 文件，其中包含 Python 扩展模块的代码。
  • .c 文件由 C 编译器编译为 .so 文件(或 Windows 上的 .pyd)，该文件可以直接导入到 Python 会话中。 setuptools 负责这一部分。尽管 Cython 可以在某些情况下为您调用它们。

下面看一个简单的样例。假设 hello.pyx 文件中有一个简单的 “hello world”脚本：

import cython
def say_hello_to(name):
    print("Hello %s!" % name)

def f(double x):
    return x ** 2 - x

def integrate_f(double a, double b, int N):
    cdef int i
    cdef double s
    cdef double dx
    s = 0
    dx = (b - a) / N
    for i in range(N):
        s += f(a + i * dx)
    return s * dx

下面是相应的 setup.py 脚本

from setuptools import setup
from Cython.Build import cythonize

setup(
    name='Hello world app',
    ext_modules=cythonize("hello.pyx"),
    zip_safe=False,
)

要构建，请运行 python setup.py build ext --inplace​ 命令。然后简单地启动一个 Python 终端，既可导入使用相关函数：

from hello import say_hello_to, integrate_f

我们也可以借助 Python 的类型注释，在写法时候更接近 Python 一点：

import cython

def f(x: cython.double):
    return x ** 2 - x


def integrate_f(a: cython.double, b: cython.double, N: cython.int):
    i: cython.int
    s: cython.double
    dx: cython.double
    s = 0
    dx = (b - a) / N
    for i in range(N):
        s += f(a + i * dx)
    return s * dx

Numba

Numba 是一个开源的、支持 numpy 的 Python 优化编译器，由 Anaconda, Inc.赞助。它使用 LLVM 从 Python 语法生成机器代码。它在使用 numpy 数组、函数和循环的代码上工作得最好。往往在数据分析、科学计算领域使用的比较多。 使用上比较简单，就是在我们 Python 函数上添加装饰器，相比其他方案（替换解释器、用其他语言编译等）工作流上侵入性比较小。

from numba import njit
import random

@njit
def monte_carlo_pi(nsamples):
    acc = 0
    for i in range(nsamples):
        x = random.random()
        y = random.random()
        if (x ** 2 + y ** 2) < 1.0:
            acc += 1
    return 4.0 * acc / nsamples

对于 numpy：

@njit(parallel=True)
def logistic_regression(Y, X, w, iterations):
    for i in range(iterations):
        w -= np.dot(((1.0 /
              (1.0 + np.exp(-Y * np.dot(X, w)))
              - 1.0) * Y), X)
    return w

Codon

Codon 是一个使用 LLVM 的高性能的 Python 编译器，它将 Python 代码编译为本机机器代码，没有任何运行时开销。相比 Python 在单个线程上加速是 10-100 倍或更多。与 Python 不同，Codon 支持本地多线程，这可以使速度提高许多倍。从生物信息领域孵化出来的项目，当前还比较新。 使用上，提供了多种方式。

from time import time

def fib(n):
    return n if n < 2 else fib(n - 1) + fib(n - 2)

t0 = time()
ans = fib(40)
t1 = time()
print(f'Computed fib(40) = {ans} in {t1 - t0} seconds.')

运行上需要使用 codon

$ python3 fib.py
Computed fib(40) = 102334155 in 17.979357957839966 seconds.
$ codon run -release fib.py
Computed fib(40) = 102334155 in 0.275645 seconds.

Codon 也提供了类似 Numba 的方案，首先安装 codon-jit​：

pip install codon-jit

代码中使用 @codon.jit ​装饰器：

import codon
from time import time

def is_prime_python(n):
    if n <= 1:
        return False
    for i in range(2, n):
        if n % i == 0:
            return False
    return True

@codon.jit
def is_prime_codon(n):
    if n <= 1:
        return False
    for i in range(2, n):
        if n % i == 0:
            return False
    return True

t0 = time()
ans = sum(1 for i in range(100000, 200000) if is_prime_python(i))
t1 = time()
print(f'[python] {ans} | took {t1 - t0} seconds')

t0 = time()
ans = sum(1 for i in range(100000, 200000) if is_prime_codon(i))
t1 = time()
print(f'[codon]  {ans} | took {t1 - t0} seconds')

性能提升也非常明显：

[python] 8392 | took 39.6610209941864 seconds
[codon]  8392 | took 0.998633861541748 seconds

Taichi

Taichi 是一种嵌在 Python 中的并行编程语言，使用 Python 语言作为 DSL，所以我们可以在正常的 Python 代码中使用它。它可以帮助我们轻松编写可移植的高性能并行程序，专注于高性能计算和图形领域。安装使用上和 Numba 类似。

pip install taichi

import taichi as ti
import taichi.math as tm

ti.init(arch=ti.gpu)

n = 320
pixels = ti.field(dtype=float, shape=(n * 2, n))

@ti.func
def complex_sqr(z):  # complex square of a 2D vector
    return tm.vec2(z[0] * z[0] - z[1] * z[1], 2 * z[0] * z[1])

@ti.kernel
def paint(t: float):
    for i, j in pixels:  # Parallelized over all pixels
        c = tm.vec2(-0.8, tm.cos(t) * 0.2)
        z = tm.vec2(i / n - 1, j / n - 0.5) * 2
        iterations = 0
        while z.norm() < 20 and iterations < 50:
            z = complex_sqr(z) + c
            iterations += 1
        pixels[i, j] = 1 - iterations * 0.02

gui = ti.GUI("Julia Set", res=(n * 2, n))

for i in range(1000000):
    paint(i * 0.03)
    gui.set_image(pixels)
    gui.show()

pybind11

pybin11 是一个轻量级的仅头文件库，它在 Python 中暴露 C++ 类型，反之亦然，主要用于创建现有 C++ 代码的 Python 绑定，它的目标和语法类似于 Boost.Python，相对大而全的 Boost 更加轻量。可是使用 C++ 来编写 Python 包。风格上大体如下：

#include <pybind11/pybind11.h>

#define STRINGIFY(x) #x
#define MACRO_STRINGIFY(x) STRINGIFY(x)

int add(int i, int j) {
    return i + j;
}

namespace py = pybind11;

PYBIND11_MODULE(python_example, m) {
    m.doc() = R"pbdoc(
        Pybind11 example plugin
        -----------------------

        .. currentmodule:: python_example

        .. autosummary::
           :toctree: _generate

           add
           subtract
    )pbdoc";

    m.def("add", &add, R"pbdoc(
        Add two numbers

        Some other explanation about the add function.
    )pbdoc");

    m.def("subtract", [](int i, int j) { return i - j; }, R"pbdoc(
        Subtract two numbers

        Some other explanation about the subtract function.
    )pbdoc");

#ifdef VERSION_INFO
    m.attr("__version__") = MACRO_STRINGIFY(VERSION_INFO);
#else
    m.attr("__version__") = "dev";
#endif
}

构建上可以参考官方的样例 GitHub - pybind/python_example: Example pybind11 module built with a Python-based build system

PyO3

PyO3 是 Python 的 Rust 绑定，包括用于创建本地 Python 扩展模块的工具。还支持从 Rust 二进制文件运行 Python 代码并与之交互。相比 pybind11，受限于 C++，Rust 工具链相对更完善一点，打包上有 maturin 方案，体验更好。 首先：

# (replace string_sum with the desired package name)
$ mkdir string_sum
$ cd string_sum
$ python -m venv .env
$ source .env/bin/activate
$ pip install maturin

使用 maturin ​创建样板项目：

maturin init

这个命令会生成的最重要的文件是 Cargo.toml 和 lib.rs，大致像下面这样

[package]
name = "string_sum"
version = "0.1.0"
edition = "2021"

[lib]
# The name of the native library. This is the name which will be used in Python to import the
# library (i.e. `import string_sum`). If you change this, you must also change the name of the
# `#[pymodule]` in `src/lib.rs`.
name = "string_sum"
# "cdylib" is necessary to produce a shared library for Python to import from.
#
# Downstream Rust code (including code in `bin/`, `examples/`, and `tests/`) will not be able
# to `use string_sum;` unless the "rlib" or "lib" crate type is also included, e.g.:
# crate-type = ["cdylib", "rlib"]
crate-type = ["cdylib"]

[dependencies]
pyo3 = { version = "0.17.3", features = ["extension-module"] }

​src/lib.rs​

// 
use pyo3::prelude::*;

/// Formats the sum of two numbers as string.
#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}

/// A Python module implemented in Rust. The name of this function must match
/// the `lib.name` setting in the `Cargo.toml`, else Python will not be able to
/// import the module.
#[pymodule]
fn string_sum(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    Ok(())
}

最后运行 maturin develop 就可以在 python 中调用了：

$ maturin develop
# lots of progress output as maturin runs the compilation...
$ python
>>> import string_sum
>>> string_sum.sum_as_string(5, 20)
'25'

其他

最近也看到字节开源的 MatxScript，是一个高性能可扩展的 Python 编译器，可以自动化把 Python 类或函数翻译成 C++，运行时完全没有 Python 开销。目前 Matx 主要支持机器学习相关的应用，但我们预期未来也会将 matx 推广到更多对灵活性和性能敏感的业务中。感兴趣的可以了解下。

总结

对于一般通用场景用户，对性能没有那么强烈的诉求，紧跟官方步伐，升级到最新版本的 Python 既可，或者使用 PyPy。Numba、Codon、Taichi 等这一类，原理上基本相同，也都是从某个场景发展过来的较为通用的方案，但是也都有各自的一些限制，按照场景选择自己顺手的即可。Cython、Pybind11、PyO3 可以看做一类，都是需要学习新语言，相对 Cython 在语言学习成本上最低，但是也更小众，Pybind11、PyO3 比较适合的场景是为现有的 C++、Rust 代码，提供 Python API，如果不想写 C++、Rust，又对性能有很高的诉求，Cython 是个不错的选择，如果是新开项目，可以尝试下逐渐受关注的 Rust。

avatar
文章
阅读
粉丝