给 Python 算法插上性能的翅膀——pybind11 落地实践

最新推荐文章于 2024-09-04 15:08:09 发布

Python_xiaowu

最新推荐文章于 2024-09-04 15:08:09 发布

阅读量2.2k

点赞数 2

文章标签： python ai pybind11 算法

本文链接：https://blog.csdn.net/Python_xiaowu/article/details/121983180

版权

本文介绍了在AI算法开发中，如何使用pybind11提升Python性能。对比了原生方案、Cython、SWIG、Boost.Python等方法，重点阐述了pybind11的优势，如轻量级、良好的C++兼容性和自动类型转换。通过C++扩展Python，利用pybind11在多线程中释放GIL锁，实现多核并行计算，显著提高性能。同时，文章提供了C++调用Python、数据类型转换的示例，并分享了在广告多媒体AI领域的应用案例。

摘要由CSDN通过智能技术生成

1. 背景

目前 AI 算法开发特别是训练基本都以 Python 为主，主流的 AI 计算框架如 TensorFlow、PyTorch 等都提供了丰富的 Python 接口。有句话说得好，人生苦短，我用 Python。但由于 Python 属于动态语言，解释执行并缺少成熟的 JIT 方案，计算密集型场景多核并发受限等原因，很难直接满足较高性能要求的实时服务需求。在一些对性能要求高的场景下，还是需要使用 C/C++来解决。但是如果要求算法同学全部使用 C++来开发线上推理服务，成本又非常高，导致开发效率和资源浪费。因此，如果有轻便的方法能将 Python 和部分 C++编写的核心代码结合起来，就能达到既保证开发效率又保证服务性能的效果。本文主要介绍 pybind11 在腾讯广告多媒体 AI Python 算法的加速实践，以及过程中的一些经验总结。

2. 业内方案

2.1 原生方案

Python 官方提供了 Python/C API，可以实现「用 C 语言编写 Python 库」，先上一段代码感受一下：

<span style="color:#444444"><span style="background-color:#f6f6f6"><span style="color:#333333"><strong>static</strong></span> PyObject *
<span style="color:#880000"><strong>spam_system</strong></span>(PyObject *self, PyObject *args)
{
    <span style="color:#333333"><strong>const</strong></span> <span style="color:#333333"><strong>char</strong></span> *command;
    <span style="color:#333333"><strong>int</strong></span> sts;

    <span style="color:#333333"><strong>if</strong></span> (!PyArg_ParseTuple(args, <span style="color:#880000">"s"</span>, &command))
        <span style="color:#333333"><strong>return</strong></span> <span style="color:#78a960">NULL</span>;
    sts = system(command);
    <span style="color:#333333"><strong>return</strong></span> PyLong_FromLong(sts);
}</span></span>

可见改造成本非常高，所有的基本类型都必须手动改为 CPython 解释器封装的绑定类型。由此不难理解，为何 Python 官网也建议大家使用第三方解决方案[1]。

2.2 赛通

Cython 主要打通的是 Python 和 C，方便为 Python 编写 C 扩展。Cython 的编译器支持转化 Python 代码为 C 代码，这些 C 代码可以调用 Python/C 的 API。从本质上来说，Cython 就是包含 C 数据类型的 Python。目前 Python 的 numpy，以及我厂的 tRPC-Python 框架有所应用。

缺点：

需要手动植入 Cython 自带语法（cdef 等），移植和复用成本高
需要增加其他文件，如 setup.py、*.pyx 来让你的 Python 代码最后能够转成性能较高的 C 代码
对于 C++的支持程度存疑

2.3 SIWG

SIWG 主要解决其他高级语言与 C 和 C++语言交互的问题，支持十几种编程语言，包括常见的 java、C#、javascript、Python 等。使用时需要用*.i 文件定义接口，然后用工具生成跨语言交互代码。但由于支持的语言众多，因此在 Python 端性能表现不是太好。

值得一提的是，TensorFlow 早期也是使用 SWIG 来封装 Python 接口，正式由于 SIWG 存在性能不够好、构建复杂、绑定代码晦涩难读等问题，TensorFlow 已于 2019 年将 SIWG 切换为 pybind11[2]。

2.4 Boost.Python

C++中广泛应用的 Boost 开源库，也提供了 Python binding 功能。使用上，通过宏定义和元编程来简化 Python 的 API 调用。但最大的缺点是需要依赖庞大的 Boost 库，编译和依赖关系包袱重，只用于解决 Python 绑定的话有一种高射炮打蚊子的既视感。

2.5 pybind11

可以理解为以 Boost.Python 为蓝本，仅提供 Python & C++ binding 功能的精简版，相对于 Boost.Python 在 binary size 以及编译速度上有不少优势。对 C++支持非常好，基于 C++11 应用了各种新特性，也许 pybind11 的后缀 11 就是出于这个原因。

Pybind11 通过 C++ 编译时的自省来推断类型信息，来最大程度地减少传统拓展 Python 模块时繁杂的样板代码，且实现了常见数据类型，如 STL 数据结构、智能指针、类、函数重载、实例方法等到 Python 的自动转换，其中函数可以接收和返回自定义数据类型的值、指针或引用。

特点：

轻量且功能单一，聚焦于提供 C++ & Python binding，交互代码简洁
对常见的 C++数据类型如 STL、Python 库如 numpy 等兼容很好，无人工转换成本
only header 方式，无需额外代码生成，编译期即可完成绑定关系建立，减小 binary 大小
支持 C++新特性，对 C++的重载、继承，debug 方式便捷易用
完善的官方文档支持，应用于多个知名开源项目

"说话很便宜，给我看看你的代码。

<span style="color:#444444"><span style="background-color:#f6f6f6">PYBIND11_MODULE (libcppex, m) {
    m.def(<span style="color:#880000">"add"</span>, [](<span style="color:#333333"><strong>int</strong></span> a, <span style="color:#333333"><strong>int</strong></span> b) -> <span style="color:#333333"><strong>int</strong></span> { <span style="color:#333333"><strong>return</strong></span> a + b; });
}</span></span>

3. Python 调C++

3.1 从 GIL 锁说起

GIL（Global Interpreter Lock）全局解释器锁：同一时刻在一个进程只允许一个线程使用解释器，导致多线程无法真正用到多核。由于持有锁的线程在执行到 I/O 密集函数等一些等待操作时会自动释放 GIL 锁，所以对于 I/O 密集型服务来说，多线程是有效果的。但对于 CPU 密集型操作，由于每次只能有一个线程真正执行计算，对性能的影响可想而知。

这里必须说明的是，GIL 并不是 Python 本身的缺陷，而是目前 Python 默认使用的 CPython 解析器引入的线程安全保护锁。我们一般说 Python 存在 GIL 锁，其实只针对于 CPython 解释器。那么如果我们能想办法避开 GIL 锁，是不是就能有很不错的加速效果？答案是肯定的，一种方案是改为使用其他解释器如 pypy 等，但对于成熟的 C 扩展库兼容不够好，维护成本高。另一种方案，就是通过 C/C++扩展来封装计算密集部分代码，并在执行时移除 GIL 锁。

3.2 Python 算法性能优化

pybind11 就提供了在 C++端手动释放 GIL 锁的接口，因此，我们只需要将密集计算的部分代码，改造成 C++代码，并在执行前后分别释放/获取 GIL 锁，Python 算法的多核计算能力就被解锁了。当然，除了显示调用接口释放 GIL 锁的方法之外，也可以在 C++内部将计算密集型代码切换到其他 C++线程异步执行，也同样可以规避 GIL 锁利用多核。

下面以 100 万次城市间球面距离计算为例，对比 C++扩展前后性能差异：

C++端：

<span style="color:#444444"><span style="background-color:#f6f6f6"><span style="color:#1f7199">#<span style="color:#333333"><strong>include</strong></span> <span style="color:#4d99bf"><math.h></span></span>
<span style="color:#1f7199">#<span style="color:#333333"><strong>include</strong></span> <span style="color:#4d99bf"><stdio.h></span></span>
<span style="color:#1f7199">#<span style="color:#333333"><strong>include</strong></span> <span style="color:#4d99bf"><time.h></span></span>
<span style="color:#1f7199">#<span style="color:#333333"><strong>include</strong></span> <span style="color:#4d99bf"><pybind11/embed.h></span></span>


<span style="color:#333333"><strong>namespace</strong></span> py = pybind11;

<span style="color:#333333"><strong>