文 @ 小 P 维修员
0 前言
深度学习算法研究如火如荼，为了更好的算法表现，算法研究者们设计出了更加复杂和精巧的算法。
主流的深度学习框架提供了灵活直观的算法表达方式和丰富的操作接口，使得研究者们可以方便地实现这些算法逻辑。但是另一方面，灵活的表达为实际的执行带来了挑战——算法中过于动态的行为会使得框架执行的效率降低。
比如检测任务中常见的检测框编解码算法，混合了灵活的索引操作和简单的算术操作。这些操作本身的计算量不大，但由于其逻辑的灵活性，使得计算不能被高效地调度和执行。
如何兼顾灵活动态的算法表达和高效的计算执行，是深度学习框架要解决的一个核心问题。
本文介绍一种以函数为单位的即时翻译技术，在保持整个算法灵活表达的基础上，对特定函数进行一定的静态化，以提升其执行效率。
正文分为三个章节，分别介绍即时翻译技术的背景、原理和挑战。
- 即时翻译技术的背景
- 即时翻译技术的原理
- 即时翻译技术的挑战
1 即时翻译技术的背景
深度学习框架提供高级语言（以下以 Python 为例）编程接口，用户使用高级语言定义数学模型，由框架执行模型训练的计算过程。
根据模型定义和模型执行的发生时间不同，深度学习框架的工作模式可分为立即执行模式和图执行模式。
- Eager Execution（立即执行模式， 简称立即模式）：高级语言调用某个算子或接口时，框架立即执行其对应的实现，因此代码在定义模型结构的同时执行计算。
- Graph Execution（图执行模式，简称图模式）：高级语言预先定义完整的模型，并由深度学习框架将模型记录为中间表示，在运行时执行中间表示对应的计算。
在图模式中，中间表示可以是一个描述计算之间拓扑关系的有向无环图，我们称之为计算图，这个图在定义之后一直存在且在每次计算都不发生改变。而在立即模式中，这个图可以是每次都变化的，或者根本不存在一个完整具体的图表示。
深度学习框架的两种工作模式各有优劣。
一方面，立即模式编写灵活，可以用纯 Python 来表达分支、循环等控制流，并且支持直接输出中间结果，便于调试；而图模式则需要使用框架自定义的接口来表达控制流逻辑，定义复杂模型的过程十分繁琐。
另一方面，图执行模式只需要生成一次计算图，易于对计算图进行图优化，在后续调用模型计算时，可以直接基于优化后的计算图进行计算，因此执行效率很高；而立即执行模式在每次调用模型时，都会重新解释 Python 模型代码并定义一个新的计算图，由于算子是逐个定义并执行的，且无法利用算子融合等图优化技术，因此执行效率低于图模式。
import tensorflow.compat.v1 as tf # TensorFlow 1.0风格的计算代码，采用图执行模式
g = tf.Graph() # 初始化计算图
with g.as_default(): # 设置为默认计算图
a = tf.constant([[2., 0.], [0., 2.]])
x = tf.placeholder(dtype=tf.float32, shape=(2, 2)) # 输入参数的占位符
b = tf.constant(1.)
y = tf.matmul(a, x) + b
init_op = tf.global_variables_initializer() # 待执行节点
with tf.Session(graph=g) as sess: # 配置会话
sess.run(init_op) # 执行节点
print(sess.run(y, feed_dict={x: [[1, 2], [3, 4]]})) # 输出结果
import tensorflow as tf # TensorFlow 2.0风格的计算代码，采用立即执行模式
def f(x): # 通过函数定义模型
a = tf.constant([[2., 0.], [0., 2.]])
b = tf.constant(1.)
y = tf.matmul(a, x) + b
return y
print(f([[1, 2], [3, 4]])) # 输出结果
如何将图执行模式计算效率高和立即执行模式便于开发与调试的优点相结合，是深度学习框架研发的重要研究方向。
诸如TensorFlow，PyTorch等主流的深度学习框架都引入了即时（Just-In-Time）翻译技术，在模型动态执行的过程中记录翻译得到的中间表示。
一种思路是将模型动态执行时生成的中间表示按参数签名缓存起来，在以后模型再被具有相同签名的参数调用时，直接使用被缓存的中间表示，从而可以省去解释Python代码的时间，同时还可以利用计算图优化或其他静态优化的方法对被缓存的中间表示做进一步的优化。
2 即时翻译技术的原理
2.1 程序语言的即时编译
即时 (Just-In-Time) 常被用在程序语言的编译实现中，因此其英文缩写 JIT 往往被翻译成即时编译。
因为解释器需要在代码执行的同时进行解释工作，相对于执行编译器编译生成代码的编译型语言，解释型语言的运行效率相对更低。为了加快程序运行速度，可以通过即时编译技术，将热点代码——也就是频繁执行的某些方法（或函数、循环体）提交给即时编译器进行编译并存储起来。在解释器进入方法入口时，会检查是否有已经编译好的方法代码，如果有，则可以直接执行编译器生成的低层代码，而不用再进行耗时较长的解释操作。
以 Java 为例，Java 虚拟机 JVM 可以通过即时编译技术来加速 Java 程序的执行。
要执行 Java 源代码，需要先通过编译器将其编译成平台无关的 Java 字节码（.class 文件），再由虚拟机 JVM 加载字节码文件并进行解释和执行。
为了提高字节码的运行速度，JVM 可以通过热点探测，检测出被频繁调用的某些方法，将它们的字节码提交给即时编译器编译成机器代码。对于执行频率较低的代码，通过解释器解释执行可以省去即时编译器的编译时间；而对于频繁执行的热点代码，即时编译技术可以显著提高其触发编译之后的代码运行速度。
2.2 深度学习模型的即时翻译
在采用立即执行模式的深度学习框架中，也可以采用类似于即时编译技术的思路。将模型动态调用过程中涉及到的热点代码通过即时翻译技术翻译成某种中间表示，在此后执行这些热点代码时，就可以直接执行中间表示对应的计算，从而省下 Python 解释器解释执行的时间，加速模型运行，我们将其称之为即时翻译。
利用即时翻译技术，可以解决动态灵活的模型定义和高效的模型执行之间的矛盾，使得用户既可以方便直观地定义模型并进行调试，又可以获得图优化带来的性能提升。
基础原理
即时翻译的操作单元通常是高级语言的函数。一般来说，深度学习框架可以将模型的整个计算过程放入一个规模较大的函数之中，也可以将模型计算的各个步骤划分成规模较小的函数，再通过嵌套调用来执行。用户可以利用 Python 装饰器，手动标记哪些函数是希望深度学习框架通过即时翻译技术来优化的。
如果一个函数只由顺序执行的数值计算构成，那么它只会被翻译成一个与之完全对应的静态计算图。但是，如果函数中包含了分支、循环等控制流逻辑，由于即时翻译的计算图是动态生成并存储的，此时，函数在动态执行过程中对应的计算图就可能不唯一，具体的计算图与动态执行过程中的判断语句逻辑有关。因此，即时翻译得到的静态计算图也可能是不唯一的。而在函数执行过程中，如果不考虑全局变量的引入，分支和循环的条件判断又总是直接或间接地由输入的实际参数决定。因此，一个函数在执行过程中生成的动态计算图，是由函数体和函数的输入参数共同决定的。
一种典型的将函数执行过程翻译为计算图的方法是跟踪 (trace)。 对于一组特定的输入参数，调用函数发生一次完整的执行，并按顺序记录其中发生的全部计算操作 (operator) 以及它们的操作数 (operand)，也就是所涉及到的变量和常量。因此，可以基于被记录的操作序列生成计算图。
跟踪的记录操作通常通过运算符重载来实现，其中发生的全部计算操作，由于 Python 的 if、while 等控制流语句不能重载，因此，生成的计算图也不会包含此类控制流逻辑。
也可以通过源码转换 (source to source) 的方法将函数代码翻译为计算图。
缓存与参数签名
如果不存储即时翻译生成的计算图，而在每次调用时都进行转换，那么其开销和立即执行模式是相近的。
通过引入缓存机制，可以将之前调用时即时翻译生成的计算图在图优化之后缓存，从而减少发生的转换次数。为了能正确对应调用函数时输入的参数组合和相应的计算图，需要按照生成时函数实际参数的参数签名，对计算图进行储存。一般来说，深度学习模型的输入参数包括张量和常量，张量通常只参与数值计算，而常量可能参与数值计算，也可能用于模型的逻辑控制。
对于输入的张量，一般来源于数据集，因此具体数值各不相同。如果在计算参数签名时按照张量里的数值储存，就会导致几乎每一次新的计算调用都触发一次即时翻译，消耗大量时间。一般情况下，输入张量的具体数值并不会影响发生的操作序列，但是其形状和类型可能会影响框架调用底层运算的具体实现，从而影响通过追踪得到的操作序列。因此，在计算张量的签名时，可以按照它的形状和所储存的数值类型（一般称为 dtype）进行记录。得到的参数签名是张量的抽象类型，代表了一类形状和数值类型均相同的张量。在实际操作中，也可以根据需求在签名中记录更多的张量属性。
而对于输入的常量，例如 int 和 float 数值，由于它们可能作为超参数影响发生的操作序列，因此，一般将它们按具体数值储存到签名之中。对于其他输入的自定义类的实例，可以按实例的id存储，但是需要注意的是，如果函数中使用了参数中类实例的属性，而不采取额外手段追踪这些对象的属性，当它们发生突变 (mutate) 时，可能会影响函数的执行逻辑，改变函数的计算序列，造成计算的错误。
可变形状支持
在实际应用过程中，有时模型会需要支持形状不固定的输入张量。例如在计算机视觉任务中，可能会遇到大小不一的图像输入。在大部分情况下，输入张量的大小并不会影响模型执行过程中发生的操作序列。举例来说，两个大小一致的张量相加，所记录的加法算子操作，并不会因为加数张量的大小变化而变化。
因此，在这类情况下，尽管输入的张量大小变化了，模型计算的计算操作序列是相同的。当输入了具有新的大小的张量，也可以沿用之前缓存的计算图，而无需通过即时翻译生成新的计算图。框架可以自动判断参数里的张量中哪些维度大小是可变的，也可以由用户手动指定哪些维度是可变维度。
3 即时翻译技术的挑战
深度学习框架的即时翻译技术还处于初期发展的阶段，面临着诸多挑战。例如上文所提到的，参数中自定义类对象的属性发生突变，可能会间接影响函数调用的操作序列，但仅仅基于重载追踪得到的计算图难以检测到这些突变。
又例如，在运行时，如果可以将控制流逻辑完整地记录成中间表达，就可以令即时翻译接受更多的 Python 原生语句，提高框架的表达能力。
深度学习框架工作者正追求以简洁而高效的方式解决这些问题，从而在未来的实际生产中，进一步应用兼具高开发效率和高执行效率的即时翻译技术。
感谢阅读，欢迎在评论区留言讨论哦~
P.S. 如果喜欢本篇文章，请多多 点赞，让更多的人看见我们 :D
关注 公众号「SenseParrots」，获取人工智能框架前沿业界动态与技术思考。