《机器学习系统：设计和实现》读后感和一些思考

Posted 2023-03-04 沉迷单车的追风少年

目录

计算图、编译器前端、编译器后端

计算图

计算图的作用

计算图的组成

静态计算图与动态计算图

编译器前端 

IR中间表示

机器学习框架的中间表示

常见编译器前端优化方法

编译器后端

概述

通用硬件优化：算子拆分和算子融合

算子信息

数据精度和存储方法

算子选择的过程

In-Place算子

模型推理

汇编语言优化

寄存器与NEON指令

PTQ训练后量化

量化公式

具体流程

---

计算图、编译器前端、编译器后端

• 计算图： 利用不同编程接口实现的机器学习程序需要共享一个运行后端。实现这一后端的关键技术是：应用无关的计算图。计算图包含计算节点，节点之间的边表达计算依赖。计算图可以被同步和异步执行。其实就是一种模型高级的中间表示。

• 编译器前端： 给定一个计算图，机器学习框架会对计算图做一系列优化。和硬件无关的优化由编译器前端实现。编译器前端实现包括：中间表达，自动微分，类型推导和静态分析等等。

• 编译器后端和运行时： 机器学习框架利用编译器后端对计算图可以进一步针对硬件的特性（例如说，L2/L3大小，指令流水线长度）进行性能优化。最终优化后的计算图通过运行时执行在通用处理器（CPU）或者是硬件加速器之上。运行时需要实现算子选择和内存分配等技术。

现代机器学习系统需要兼有易用性和高性能，因此其一般选择Python作为前端编程语言，而使用C和C++作为后端编程语言。前端负责静态分析、类型推导以及自动微分、分布式并行子图拆分等PASS优化；后端负责硬件相关的优化，如：内存优化、图算融合等。

计算图

计算图的作用

• 对于输入数据、算子和算子执行顺序的统一表达。 机器学习框架用户可以用多种高层次编程语言（Python，Julia和C++）来编写训练程序。这些高层次程序需要统一的表达成框架底层C和C++算子的执行。因此，计算图的第一个核心作用是可以作为一个统一的数据结构来表达用户用不同语言编写的训练程序。这个数据结构可以准确表述用户的输入数据、模型所带有的多个算子，以及算子之间的执行顺序。

• 定义中间状态和模型状态。 在一个用户训练程序中，用户会生成中间变量（神经网络层之间传递的激活值和梯度）来完成复杂的训练过程。而这其中，只有模型参数需要最后持久化，从而为后续的模型推理做准备。通过计算图，机器学习框架可以准确分析出中间状态的生命周期（一个中间变量何时生成，以及何时销毁），从而帮助框架更好的管理内存。

• 自动化计算梯度。 用户给定的训练程序仅仅包含了一个机器学习模型如何将用户输入（一般为训练数据）转化为输出（一般为损失函数）的过程。而为了训练这个模型，机器学习框架需要分析任意机器学习模型和其中的算子，找出自动化计算梯度的方法。计算图的出现让自动化分析模型定义和自动化计算梯度成为可能。

• 优化程序执行。 用户给定的模型程序往往是“串行化”地连接起来多个神经网络层。通过利用计算图来分析模型中算子的执行关系，机器学习框架可以更好地发现将算子进行异步执行的机会，从而以更快的速度完成模型程序的执行。

计算图的组成

计算图由基本数据结构：张量(Tensor)和基本运算单元算子(Operator)构成。

静态计算图与动态计算图

1、静态计算图

静态计算图采用先编译后执行的方式，该模式将计算图的定义和执行进行分离。 

在静态图模式下使用前端语言定义模型形成完整的程序表达后，并不使用前端语言解释器进行执行，而是将前端描述的完整模型交给计算框架。框架在执行模型计算之前会首先对神经网络模型进行分析，获取网络层之间的连接拓扑关系以及参数变量设置、损失函数等信息，接着用一种特殊的静态数据结构来描述拓扑结构及其他神经网络模型组件，这种特殊的静态数据结构通常被称为静态计算图。静态计算图可以通过优化策略转换成等价的更加高效的结构。当进行模型训练或者推理过程时，静态计算图接收数据并通过相应硬件调度执行图中的算子来完成任务。

2、动态计算图 

动态计算图采用解析式的执行方式，其核心特点是编译与执行同时发生。

动态图采用前端语言自身的解释器对代码进行解析，利用计算框架本身的算子分发功能，算子会即刻执行并输出结果。动态图模式采用用户友好的命令式编程范式，使用前端语言构建神经网络模型更加简洁。

3、两者对比

静态生成和动态生成的过程各有利弊。从使用者的角度可以直观的感受到静态图不能实时获取中间结果、代码调试困难以及控制流编写复杂，而动态图可以实时获取结果、调试简单、控制流符合编程习惯。虽然静态图的编写、生成过程复杂，但是相应的执行性能却超过动态图。

编译器前端 

IR中间表示

中间表示(IR)，是编译器用于表示源代码的数据结构或代码，是程序编译过程中介于源语言和目标语言之间的程序表示。几乎所有的编译器都需要某种形式的中间表示，来对被分析、转换和优化的代码进行建模。在编译过程中，中间表示必须具备足够的表达力，在不丢失信息的情况下准确表达源代码，并且充分考虑从源代码到目标代码编译的完备性、编译优化的易用性和性能。

在此基础上，编译流程就可以在前后端直接增加更多的优化流程，这些优化流程以现有IR为输入，又以新生成的IR为输出，被称为优化器。优化器负责分析并改进中间表示，极大程度的提高了编译流程的可拓展性，也降低了优化流程对前端和后端的破坏。

机器学习框架的中间表示

在设计机器学习框架的中间表示时，需要充分考虑以下因素：

1) 张量表达。机器学习框架主要处理张量数据，因此正确处理张量数据类型是机器学习框架中间表示的基本要求。

2) 自动微分。自动微分是指对网络模型的自动求导，通过梯度指导对网络权重的优化。主流机器学习框架都提供了自动微分的功能，在设计中间表示时需要考虑自动微分实现的简洁性、性能以及高阶微分的扩展能力。

3) 计算图模式。主流机器学习框架如TensorFlow、PyTorch、MindSpore等都提供了静态图和动态图两种计算图模式，静态计算图模式先创建定义计算图，再显式执行，有利于对计算图进行优化，高效但不灵活。动态计算图模式则是每使用一个算子后，该算子会在计算图中立即执行得到结果，使用灵活、便于调试，但运行速度较低。机器学习框架的中间表示设计同时支持静态图和动态图，可以针对待解决的任务需求，选择合适的模式构建算法模型。

4) 支持高阶函数和闭包 。高阶函数和闭包是函数式编程的重要特性，高阶函数是指使用其它函数作为参数、或者返回一个函数作为结果的函数，闭包是指代码块和作用域环境的结合，可以在另一个作用域中调用一个函数的内部函数，并访问到该函数作用域中的成员。支持高阶函数和闭包，可以抽象通用问题、减少重复代码、提升框架表达的灵活性和简洁性。

5) 编译优化。机器学习框架的编译优化主要包括硬件无关的优化、硬件相关的优化、部署推理相关的优化等，这些优化都依赖于中间表示的实现。

6) JIT(Just In Time)能力。机器学习框架进行编译执行加速时，经常用到JIT即时编译。JIT编译优化将会对中间表示中的数据流图的可优化部分实施优化，包括循环展开、融合、内联等。中间表示设计是否合理，将会影响机器学习框架的JIT编译性能和程序的运行能力。

常见编译器前端优化方法

1、无用与不可达代码消除

如 图6.5.2所示。无用代码是指输出结果没有被任何其他代码所使用的代码。不可达代码是指没有有效的控制流路径包含该代码。删除无用或不可达的代码可以使得中间表示更小，提高程序的编译与执行速度。无用与不可达代码一方面有可能来自于程序编写者的编写失误，也有可能是其他编译优化所产生的结果。

2、常量传播、常量折叠

常量传播：如 图6.5.3所示，如果某些量为已知值的常量，那么可以在编译时刻将使用这些量的地方进行替换。

常量折叠：如 图6.5.3所示，多个量进行计算时，如果能够在编译时刻直接计算出其结果，那么变量将由常量替换。

3、公共子表达式消除

如 图6.5.4所示，如果一个表达式E已经计算过了，并且从先前的计算到现在E中所有变量的值都没有发生变化，那么E就成为了公共子表达式。对于这种表达式，没有必要花时间再对它进行计算，只需要直接用前面计算过的表达式结果代替E就可以了。

编译器后端

概述

如 图7.1.1所示，编译器后端处于前端和硬件驱动层中间，主要负责计算图优化、算子选择和内存分配的任务。首先，需要根据硬件设备的特性将IR图进行等价图变换，以便在硬件上能够找到对应的执行算子，该过程是计算图优化的重要步骤之一。前端IR生成是解析用户代码，属于一个较高的抽象层次，隐藏一些底层运行的细节信息，此时无法直接对应硬件上的算子（算子是设备上的基本计算序列，例如MatMul、Convolution和ReLU等），需要将细节信息进行展开后，才能映射到目标硬件上的算子。对于某些前端IR的子集来说，一个算子便能够执行对应的功能，此时可以将这些IR节点合并成为一个计算节点，该过程称之为算子融合；对于一些复杂计算，后端并没有直接与之对应的算子，但是可以通过几个基本运算的算子组合达到同样的计算效果，此时可以将前端IR节点拆分成多个小算子。然后，我们需要进行算子选择。算子选择是在得到优化的IR图后，需要选取最合适的目标设备算子。针对用户代码所产生的IR往往可以映射成多种不同的硬件算子，但是生成不同的算子执行效率往往有很大的差别，如何根据前端IR选择出最高效的算子，是算子选择的核心问题。算子选择本质上是一个模式匹配问题。其最简单的方法就是每一个IR节点对应一个目标硬件的算子，但是这种方法往往对目标硬件的资源利用比较差。目前来说对于现有的编译器一般都对每一个IR节点提供了多个候选的算子，算子选择目标就是从中选择最优的一个算子作为最终执行在设备上的算子。总的来说，在机器学习系统中，对前端生成的IR图上的各个节点进行拆分和融合，让前端所表示的高层次IR逐步转换为可以在硬件设备上执行的低层次IR。得到了这种更加贴合硬件的IR后，对于每个单节点的IR可能仍然有很多种不同的选择，例如可以选择不同的输入输出格式和数据类型，我们需要对IR图上每个节点选择出最为合适的算子，算子选择过程可以认为是针对IR图的细粒度优化过程，最终生成完整的算子序列。最后，遍历算子序列，为每个算子分配相应的输入输出内存，然后将算子加载到设备上执行计算。

通用硬件优化：算子拆分和算子融合

深度学习算子按其对资源的需求可以分为两类： 计算密集型算子，这些算子的时间绝大部分花在计算上，如卷积、全连接等； 访存密集型算子，这些算子的时间绝大部分花在访存上，他们大部分是Element-Wise算子，例如 ReLU、Element-Wise Sum等。 在典型的深度学习模型中，一般计算密集型和访存密集型算子是相伴出现的，最简单的例子是“Conv + ReLU”。Conv卷积算子是计算密集型，ReLU算子是访存密集型算子，ReLU算子可以直接取Conv算子的计算结果进行计算，因此我们可以将二者融合成一个算子来进行计算，从而减少内存访问延时和带宽压力，提高执行效率。

例如：“Conv + Conv + Sum + ReLU”的融合，从 图7.2.1中我们可以看到融合后的算子减少了两个内存的读和写的操作，优化了Conv的输出和Sum的输出的读和写的操作。

除了上述针对特定算子类型结构的融合优化外，基于自动算子生成技术，还可以实现更灵活、更极致的通用优化。以 MindSpore 的图算融合技术为例，图算融合通过“算子拆解、算子聚合、算子重建”三个主要阶段（如图）让计算图中的计算更密集，并进一步减少低效的内存访问。

图7.2.2中，算子拆解阶段（Expander）将计算图中一些复杂算子（composite op，图中Op1、Op3、Op4）展开为计算等价的基本算子组合（ 图中虚线正方形框包围着的部分）；在算子聚合阶段（Aggregation），将计算图中将基本算子（basic op，如图中Op2）、拆解后的算子（expanded op）组合融合，形成一个更大范围的算子组合；在算子重建阶段（Reconstruction）中，按照输入tensor到输出tensor的仿射关系将基本算子进行分类：elemwise、 broadcast、reduce、transform等，并在这基础上归纳出不同的通用计算规则（如 elemwise + reduce 规则：elemwise + reduce在满足一定条件后可以高效执行），根据这些计算规则不断地从这个大的算子组合上进行分析、筛选，最终重新构建成新的算子（如图中虚线正方形包围的两个算子 New Op1 和 New Op2）。图算融合通过对计算图结构的拆解和聚合，可以实现跨算子边界的联合优化；并在算子重建中，通过通用的计算规则，以必要的访存作为代价，生成对硬件更友好、执行更高效的新算子。

算子信息

1. 针对不同特点的计算平台和不同的算子，为了追求最好的性能，一般都需要选择不同的数据排布格式。机器学习系统常见的数据排布格式有NCHW和NHWC等。

2. 对于不同的硬件支持不同的计算精度，例如float32、float16和int32等。算子选择需要在所支持各种数据类型的算子中选择出用户所设定的数据类型最为相符的算子。

数据精度和存储方法

通常深度学习的系统，一般使用的是单精度float（Single Precision）浮点表示。这种数据类型占用32位内存。还有一种精度较低的数据类型float16，其内部占用了16位的内存。由于很多硬件会对float16数据类型进行优化，float16半精度的计算吞吐量可以是float32的2∼8倍，且float16可以占用的数据更小，这样可以输入更大的BatchSize，进而减少总体训练时间。接下来我们详细看一下半精度浮点数与单精度浮点数的区别。

如 图7.3.5其中sign代表符号位，占1位，表示了机器数的正负，exponent表示指数位，Mantissa为尾数位。其中float16类型的数据采用二进制的科学计数法转换为十进制的计算方式如下：

算子选择的过程

其中算子信息主要包括了支持设备类型、数据类型和数据排布格式三个方面。经过编译器前端类型推导与静态分析的阶段后，IR图中已经推导出了用户代码侧的数据类型。下面介绍算子选择的基本过程。

首先，选择算子执行的硬件设备。不同的硬件设备上，算子的实现、支持数据类型、执行效率通常会有所差别。这一步往往是用户自己指定的，若用户未指定，则编译器后端会为用户匹配一个默认的设备。 然后，后端会根据IR图中推导出的数据类型和内存排布格式选择对应的算子。

理想情况下算子选择所选择出的算子类型，应该与用户预期的类型保持一致。但是由于软硬件的限制，很可能算子的数据类型不能满足用户所期待的数据类型，此时需要对该节点进行升精度或者降精度处理才能匹配到合适的算子。

算子的数据排布格式转换是一个比较耗时的操作，为了避免频繁的格式转换所带来的内存搬运开销，数据应该尽可能地以同样的格式在算子之间传递，算子和算子的衔接要尽可能少的出现数据排布格式不一致的现象。另外，数据类型不同导致的降精度可能会使得误差变大，收敛速度变慢甚至不收敛，所以数据类型的选择也要结合具体算子分析。

In-Place算子

在内存分配流程中，我们会为每个算子的输入和输出都分配不同的内存。然而对很多算子而言，为其分配不同的输入和输出地址，会浪费内存并且影响计算性能。例如优化器算子，其计算的目的就是更新神经网络的权重；例如Python语法中的’+=‘和’*=‘操作符，将计算结果更新到符号左边的变量中；例如’a[0]=b’语法，将’a[0]’的值更新为’b’。诸如此类计算有一个特点，都是为了更新输入的值。下面以Tensor的’a[0]=b’操作为例介绍In-Place的优点。 图7.4.6左边是非In-Place操作的实现，step1将Tensor a拷贝到Tensor a’，step2将Tensor b赋值给Tensor a’，step3将Tensor a’拷贝到Tensor a。 图7.4.6右边是算子In-Place操作的实现，仅用一个步骤将Tensor b拷贝到Tensor a对于的位置上。对比两种实现，可以发现In-Place操作节省了两次拷贝的耗时，并且省去了Tensor a’内存的申请。

模型推理

汇编语言优化

对于已知功能的汇编语言程序来说，计算类指令通常是固定的，性能的瓶颈就在非计算指令上。计算机各存储设备类似于一个金字塔结构，最顶层空间最小，但是速度最快，最底层速度最慢，但是空间最大。L1-L3统称为cache(高速缓冲存储器)，CPU访问数据时，会首先访问位于CPU内部的cache，没找到再访问CPU之外的主存，此时引入了缓存命中率的概念来描述在cache中完成数据存取的占比。要想提升程序的性能，缓存命中率要尽可能的高。

下面简单列举一些提升缓存命中率、优化汇编性能的手段：

（1）循环展开：尽可能使用更多的寄存器，以代码体积换性能；

（2）指令重排：打乱不同执行单元的指令以提高流水线的利用率，提前有延迟的指令以减轻延迟，减少指令前后的数据依赖等；

（3）寄存器分块：合理分块NEON寄存器，减少寄存器空闲，增加寄存器复用；

（4）计算数据重排：尽量保证读写指令内存连续，提高缓存命中率；

（5）使用预取指令：将要使用到的数据从主存提前载入缓存，减少访问延迟。

寄存器与NEON指令

ARMv8系列的CPU上有32个NEON寄存器v0-v31，如 图10.4.2所示，NEON寄存器v0可存放128bit的数据，即4个float32，8个float16，16个int8等。

针对该处理器，可以采用SIMD(Single Instruction，Multiple Data，单指令、多数据)提升数据存取计算的速度。相比于单数据操作指令，NEON指令可以一次性操作NEON寄存器的多个数据。例如：对于浮点数的fmla指令，用法为fmla v0.4s, v1.4s, v2.4s，如 图10.4.3所示，用于将v1和v2两个寄存器中相对应的float值相乘累加到v0的值上。

PTQ训练后量化

量化公式

假设r表示量化前的浮点数，量化后的整数q可以表示为：

round(⋅)和clip(⋅)分别表示取整和截断操作，和是量化后的最小值和最大值。s是数据量化的间隔，z是表示数据偏移的偏置，z为0的量化被称为对称（Symmetric）量化，不为0的量化称为非对称（Asymmetric）量化。对称量化可以避免量化算子在推理中计算z相关的部分，降低推理时的计算复杂度；非对称量化可以根据实际数据的分布确定最小值和最小值，可以更加充分的利用量化数据信息，使得计算精度更高。 

具体流程

• 使用直方图统计的方式得到原始FP32数据的统计分布；

• 在给定的搜索空间中选取若干个和分别对激活值量化，得到量化后的数据；

• 使用直方图统计得到的统计分布;

• 计算每个与的统计分布差异，并找到差异性最低的一个对应的和来计算相应的量化参数，常见的用于度量分布差异的指标包括KL散度(Kullback-Leibler Divergence)、对称KL散度(Symmetric Kullback-Leibler Divergence)和JS散度(Jenson-Shannon Divergence)。