Python 加速好帮手：Numba

一次无可奈何的 Python 提速。

---

先简单介绍一下背景。

这门课是 Visualization，由两个作业组成。第一个作业的代码部分，要求补完一个体渲染器（volume renderer）。这个渲染器是完全软件渲染的，而且需要支持几种渲染方式（也就是几种软件着色器）：切片（代码已给出）、最大密度投影（MIP）、体元合成（compositing）。不管哪一种，核心都是 ray casting。严格来说，每一个都有“合成”步骤，不过最后一个实际上指的是最常见的、使用 Phong 模型的着色。

今年我就被一个认识我的人叫过去帮忙了。去年我修这门课的时候，用的是 Java。今年据说因为 Java 入门（一个 master level 的课程，本来在 Q1，也就是 visualization 的前一个 quartile）移到了 bachelor level 3，所以在改革之后这门课默认所有选课的人都没有 Java 基础。因此，使用的语言改为了 Python。

我不得不严重吐槽今年的老师，虽然我并不认识他。内容比较长而且跟题目切合度不算太大，所以就放在文末。

好，接着就讲优化的事。下文的渲染时间以提供的橘子的数据为准。

由于各种坑爹因素的综合作用，切片的每帧渲染时间就已经超过了2秒。至于体元合成，因为所有体元基本都要至少被投射一次，总时间可达每帧10分钟（600秒）左右。其他的已经实现了这个的小组也反馈说，用时10分钟“是非常正常的”。

正常个毛！作为读者而不是这个坑爹软件的使用者，你可能不是很清楚它对人精神的打击有多大。首先，切换到体元合成模式，渲染初始帧，10分钟。接着，你会发现默认的迁移函数（transfer function）效果很差，需要调整曲线。当你调整第一个数值的时候，新的渲染开始了，又是10分钟。在这段时间内，你对迁移函数的调整完全不会反映出来——因为上一帧还没有渲染完成，所有更改都不会应用。你根据密度直方图大致估摸出了一个曲线，但为了它能起作用，你必须在上一帧渲染完成后，再次触发更改（比如调整颜色）才能获得正确的渲染结果，于是又要等10分钟。然后，你会看到默认的视角并不能看到什么有意义的东西，于是你得尝试调整视角。跟迁移函数一样，如果正在渲染，则新的视角不会应用。这就是多个10分钟。在你获得最后的满意的图像之前，精神上要受到长时间的折磨。说“折磨”，是因为这个时间根本就是因为种种错误导致的浪费生命，而且不止是一个人的生命。

怎么办呢？我们来想想方案的必要条件：

• 使用 Python（恶心的规定）
• 非侵入式的
• 支持 GUI 实例

补充条件：

• 应用容易

其实说到底，就是加速。加速，无外乎就是提高硬件利用率、合理利用硬件。例如，并行（CPU、GPU）、更高效率的代码（设计上、编译后）、提高缓存命中等等。

各位或许听说过 Cython。但是很可惜，虽然 Python 代码是合法的 Cython 代码，但是为了让 Cython 达到最高效率，必须要将代码改写成其方言。这就违反了非侵入式的条件，代码提交后别人也不容易运行这些代码。

Python 以性能为代价提供了极高的自由，而这给解释器设计让生成高效率的机器码带来的巨大的困难。所以不可能寄希望于 Python 解释器自身给我们提供极致性能。

我也试过用 multiprocessing 模块来加速。然而这并没有什么卵用。需要优化的函数是一个虚函数，而且位于事件回调，跟 GUI 有相当的耦合。这就不像一些常用 multiprocessing 优化的模块，那些模块基本是用于数值计算。而且它们的函数是全局的（对于子进程来说是入口），参数也利于封送，运行时产生的都是后台进程。当我简单使用 multiprocessing 的时候，这个程序直接弹出了好几个窗口，把我吓了一跳——而且它们没有一个在干活。

threading？别忽悠人了。

还有一些其他的方法，不过投入的成本太高，对于一个作业，精力有限，做不起。（见文末讨论。）

---

这时候偶然搜索到了 Numba。它和 Cython 的思路是相似的，在一定条件下，转换成 C 代码后编译运行，利用 C 的速度优势（“C 的速度优势”不是很准确，但是在此不展开了），提高代码执行效率。在满足一些苛刻的条件时，甚至还可以应用并行化，进一步提高执行效率。

就使用体验而言，Numba 总体来说还是相当友好的。在正式用到作业项目之前，我用一篇文章中的示例简单测试了一下优化效率。

import time
from typing import Callable

import numpy as np
from numba import jit

def pairwise_python(X: np.ndarray) -> np.ndarray:
    M = X.shape[0]
    N = X.shape[1]
    D = np.empty((M, M), dtype=np.float)
    for i in range(M):
        for j in range(M):
            d = 0.0
            for k in range(N):
                tmp = X[i, k] - X[j, k]
                d += tmp * tmp
            D[i, j] = np.sqrt(d)
    return D

@jit(nopython=True)
def pairwise_numba(X: np.ndarray) -> np.ndarray:
    M = X.shape[0]
    N = X.shape[1]
    # https://github.com/numba/numba/issues/3993
    # Use float_/int_ instead of float{XX}/int{XX}
    D = np.empty((M, M), dtype=np.float_)
    for i in range(M):
        for j in range(M):
            d = 0.0
            for k in range(N):
                tmp = X[i, k] - X[j, k]
                d += tmp * tmp
            D[i, j] = np.sqrt(d)
    return D

def pairwise_numpy(X: np.ndarray) -> np.ndarray:
    return np.sqrt(((X[:, None, :] - X) ** 2).sum(-1))

def run_python() -> None:
    X = np.random.random((1000, 3))
    pairwise_python(X)

def run_numba() -> None:
    X = np.random.random((1000, 3))
    pairwise_numba(X)

def run_numpy() -> None:
    X = np.random.random((1000, 3))
    pairwise_numpy(X)

def time_it(name: str, func: Callable[[], None]) -> None:
    time_start = time.time()
    func()
    time_end = time.time()
    print(f"{name} used {time_end - time_start} second(s)")

if __name__ == "__main__":
    time_it("raw python", run_python)
    time_it("numba", run_numba)
    time_it("numpy", run_numpy)

在我的机器上运行，结果：

raw python used 5.295958518981934 second(s)
numba used 0.3148195743560791 second(s)
numpy used 0.05596804618835449 second(s)

可见，仅仅是转换成等价 C 代码（还没有并行）并编译，就已经获得了巨大的效率提升。当然，这还是比不过 NumPy 的向量化操作（ufunc）。

不过，在这个作业里，每个最终像素点的实际生成过程是异构（heterogeneous）的，因此无法将其向量化，自然无法充分利用 NumPy。使用 Numba 并行是可能的；考虑到作业中的这个函数远比上面的简单例子复杂，尤其是输入的参数，所以尽管理论上可能，我也没这么多精力去适配它的条件了。

我们继续来看 Numba。Numba 提供了两种模式，一种是 nopython 模式，一种是 object 模式。前者不需要涉及 Python 的解释器，后者需要。什么时候需要呢？访问除了特定受支持的数据类型（int、float、numpy.ndarray 等等）之外的，尤其是自定义的复杂类型（简单类型可以用 @jitclass）。很明显，一旦和 Python 解释器扯上关系，那么效率就得大打折扣。再加上 JIT 是有开销的，所以最后总体甚至会慢于直接在解释器中执行。因此，一般都得用 nopython 模式。

接下来就是痛苦的转换了。毕竟 Numba 要能应用，是有限制条件的。

所以第一步，就是尽量做到与 OO 无关，而是回退到面向过程的方式。毕竟文档自己说了“越像 C 越好”（笑）。拆！拆！拆！好在原程序大部分都是不需要跟对象进行交互的，因此将主体变换为面向过程的方式并没有太复杂，就是繁琐一些，花一个小时就搞定了。同时考虑到 Python 无法内联的特性，分块大小偏大，减少调用开销。

不过里面有一个问题如鲠在喉；甚至说，如果不将它解决，那么其他的优化都无法生效，Numba 必然会回退到 object 模式。在循环的最深处，需要根据迁移函数获取在该视角（viewpoint）下某个体素（voxel）的对应颜色值。这可是插在层层循环中最深的地方的一个钉子。而必需的迁移函数，在实现上是一个复杂对象，@jitclass 还不支持。嗯，仔细观察一下，颜色计算其实是根据一张表实现的，而这张表在单次渲染中不发生变化。所以我就写了一个计算量不大的函数，将迁移函数的内部状态提取到了一个 ndarray 中，并修改颜色计算使其用这个 ndarray 而不是原来的迁移函数对象。这样核心计算就可以在 nopython 模式下进行了。

目前（0.46）Numba 还不支持 MAKE_FUNCTION opcode。Numba 只会报错说“op_MAKE_FUNCTION 未实现”。这么一个莫名其妙的错误信息还是得猜一下。结果是，它不支持函数内定义的函数，无论内部的函数是否形成闭包。所以我只好又违反了模块化原则，把一些内部的小函数提出来放到全局范围里。

最后优化到了什么程度呢？原先的每帧渲染时间是大约600秒，优化后冷启动24秒（含 JIT 时间），多次运行18秒。整整33倍！而且这仅仅是将 Python 代码 JIT 的结果，甚至都没有并行化。这个函数输入太多，签名复杂，所以我就不浪费时间去做并行化了——软件本身只是用来出结果的，后面还有报告呢，这时候离截止日期已经不远了。比起再花几个小时尝试优化到最佳状态，还不如就18秒先用着，毕竟看模型又不是重复数十次上百次的重复劳动，最后总计花费时间也不会太久。

---

接下来就是纸上谈兵环节了。

我们看到了解释代码到原生代码的效率提升。如果要进一步优化异构任务的话，该怎么办呢？在这里再加一个前提：跟宿主语言无关，也就是要找一个相对通用的解。

一个很容易想到的方案是使用传统的 GPU。将计算逻辑变换成 shader，输入变换为顶点数据，输出从图像中解码。使用的 API 可以是 OpenGL（vertex/fragment）、Direct3D。这种做法别说现代 GPU，老一点的 GPU 都支持。缺点就是结果的复杂度不能太高，毕竟单元输出最多只有4个单精度浮点数。

如果使用 GPGPU，就可以避开许多输入输出的限制。可以使用的有 OpenGL（compute）、DirectCompute，以及 OpenCL 和 CUDA。其中一些还支持异构设备计算。

以上的方案，只要所使用的语言提供对应的绑定，就可以使用。

如果放宽“语言无关”限制，只针对 C/C++ 的话，有 OpenMP、OpenACC 和 SYCL。它们的效率提升方式不尽相同，有自动多线程并行的，有（合适时）使用 GPU 的。但是它们几乎可以和 C/C++ 无缝集成，而不需要将计算逻辑分离到独立的管线中。

---

对这个老师的吐槽。

客观来讲，他犯了一个严重的选型错误。他应该是这么想的，Python 不是什么“新时代的基础语言”、“万能胶水语言”，“上手十分简单”嘛。但是很明显，Python 和 VB 一样，学习曲线的上升幅度开始平缓，接着立马就极端地陡峭，也就是易于上手，难于深入，更难于精通。Python 的运行效率不能说差，但是很不擅长计算密集型的任务（至少在包括默认的 CPython 上的大部分实现，因为有 GIL）。（另外插一句，PEP-554 安排在 Python 3.9 了，所以现在也用不了。）而这门课要做的是一个软件渲染器，不仅是一个计算密集型的任务，而且是一个异构任务（不能简单地通过向量化完成）。写出逻辑是很简单的，但是一旦运行起来，就会发现它的运行时间无法忍受。想优化？对不起，优化异构任务的难度，已经远远超出了花样使用 NumPy（优化同构任务）；后者也早就不是初学者范围内的东西了。也就是说，这个作业在目标之外设置了一个巨大的障碍，而所有人不得不做出选择——要不就克服这个障碍，要不就得忍受折磨——而几乎没有人有能力克服。这是完全没有意义的精神消耗，至少对于学生而言。（现实工作中也有可能会遇到这样的问题，不过都已经到了工作了，自然就不能假设有保护了。）

他貌似对自己的水平没有准确的估计。由于很多人被折磨得死去活来，所以他似乎收到了一些问问题的邮件。将这些整理成 FAQ 之后，其中就有令人哭笑不得的回答。在程序中有一个 bug，它导致计算出的图像不会绘制到窗口中，也就是体现为“看不见”。这几乎是致命的问题，自然有人会问。而他的回答呢？“看不见图像的原因，很可能是因为你的计算机在计算上出了问题。”我真不知道该说什么好。

在代码上，可以明显看出 Java 移植的痕迹。在命名上和一些辅助类的结构上就看出来。除了 GUI 部分的代码遵循 wxPython 的一般规则、OpenGL（仅用来画包围盒）用 PyOpenGL，其他都是 Java 的简单移植。不过要说命名，Python 本身就已经很混乱了，缩写的、连写（无分隔）的、snake_case 的、camelCase 的，乱成一锅粥。以颜色的记录来说，OO 的组织方式比较利于理解，而直接使用 ndarray 更利于和 NumPy 交互，也利于优化。这更多地是一种设计选择，但是我认为有经验的话不会做出这样的选择。

程序本身交互极差。上面也提到了一部分。这个作业在最后还要完成第四种模式，将你所选定的可视化方式应用于分块的体数据。这相当于有一个一级列表，让你选择某个物体，然后再有一个二级列表，选择其在某个时间点下的数据。但是要进入这第四种模式可不是那么简单。这个模式对应的单选框默认是无效（灰色）的，你得先手工选择两个文件夹（体数据索引、各个时间），再选定了其中一个物体的其中一个时间点后，这个单选框才会激活。要不然，想切换都没法切换。这个操作本身就非常反直觉，如果没有那同学的讲解，我恐怕也得摸索个几分钟。在代码逻辑上，单物体（直接打开体数据）和多物体（本段所述的分块选择）的交互是两套代码……

由于不断地有人哀嚎渲染一帧实在是太久了，这老师在截止日期前几天发了一则公告，说如果想加速的话推荐使用 Numba。然而他在给出的例子是切片——最简单的渲染模式，基本上就是简单的 ndarray 读取，没有复杂循环，没有合成。这个直接套个 @jit(nopython=True) 就搞定了，非常简单。然而真正有优化需求的体元合成，要用上 Numba，那得知道 Numba 是怎么工作的、它对数据结构是怎么支持的、Python 本身的瓶颈在哪里、什么时候 Python 解释器是必需的（例如让人又爱又恨的 inspect），这样才能知道什么能加速什么不行，如果不行能不能换成等价的写法使得它能被加速。这在我看来，对知识厚度的要求远远超过这门课。将这种任务放在一个作业里，就是高射炮打蚊子。而且看来，他也没有对任务难度的合理估计，以为后面都只是像他给的例子那样，直接加个装饰器（decorator）那么简单。

不过，不是没有厉害的老师的。在我选过的课里面，就发生过：

• Foundations of Machine Learning。老师讲得很好，而且看代码也是简明扼要，清晰易懂，几乎没有冗余。
• Recommender Systems，和 FML 是同一个老师。在第四个作业中，我们要实现一个 VAE。其中助教提示的是“别忘了用 flatten() 调整数据的形状”，但是在查阅资料后我判断应该用 batch_flatten()，就这么写了。后来在作业的点评里这一条被专门指了出来，说确实应该用 batch_flatten()，是他失误了。
• VLSI Programming。就是我之前写过的老先生。图的构造代码写得非常漂亮，而且指出过我的一些 pythonic 的写法（例如列表推导）有效率问题——在这门课里，效率是一个很重要的衡量因素。
• Programming Methods。这门课其实是一个本科3级的课程，偏向于软件工程，我自愿选的。最后的一个作业是写一个 Takuzu（类似于数独，不过只填入0和1）的求解器，而且是可视化的。这个程序本来是用作各种设计模式的练习的，不过它的交互设计和代码水平都很不错。