我针对机器学习的 Helloword 项目 Mnist 手写数字识别，做了一个小实验，分别在我的 MacBook Pro 和 iPhone 手机上运行了同样的算法模型，把两侧的训练样本、模型结构、模型参数、训练参数等对齐，最终得出图 4-4 的结果：面对 60000 个训练样本 10 个 Epoch 在 i7 CPU 的 2015 款 15 寸 MacBook Pro 上需要 128 秒，而在我的 iPhone 13 Pro Max 上只需要 86 秒，这足以证明端上计算能力能够满足我们使用模型进行预测乃至训练模型的计算能力要求。
图 4-4 现在移动端有多快？同样的模型在 iOS 上需要 86 秒而 i7 MacOS 要 128 秒
但是，在性能一般的手机终端上，还需要测量算力要求，并针对框架和平台优化算法模型来保证用户体验。这里的算法模型优化由两部分构成，一部分源于算法模型本身的压缩、剪枝、量化、知识蒸馏等，另一部分源于框架自带的工具或 JAX、TVM 第三方工具把算法模型针对性转换成不同平台上优化的模型。第一部分可以从机器学习相关著作里学习，下面着重介绍第二部分，围绕 TensorFlow 机器学习框架提供的工具，以及 iOS 机器学习技术，分别介绍端智能的技术工程基础。
评估和准备算法模型
在开始评估和准备算法模型之前，先介绍一下前文中 iPhone 13 Pro Max 的处理器。这块儿 A15 处理器采用台积电 5nm 工艺，共集成了 150 亿个晶体管，NPU 性能达到 15.8TOPS。TOPS（Tera Operations Per Second），1TOPS 代表处理器每秒钟可进行一万亿次操作。相对硬件的算力，算法模型的复杂度相对的用算力要求 FLOPs （floating point operations）来评估，指浮点运算计算量, 可以用来衡量算法模型的复杂度。从图 4-5 中可以得知，算法模型的参数、模型大小、FLOPs 都是用来评估算法模型的指标。由于这几个指标几乎呈现参数规模越大模型越大且算力要求越高的规律，因此，在进行算法模型评估的时候，借助图 4-4 中 模型参数规模 Total params: 585958 就可以大概评估出模型的算力要求。但是，也必须注意 AlexNet 和 ResNet152 这种参数少却模型大、FLOPs 更高的情况，这主要是由于网络结构、优化器不同等因素造成的。
当然，如果要较为精确的评估来识别模型参数少却对算力要求高的情况，图 4-4 中的算力要求仅就模型拥有 58 万参数进行估算是不够的，可以用 TensorFlow 的 API 计算出该模型需要的算力为 2.8 MFLOPs 。
当然，如果使用 Pytorch 算法框架，则需要对应的工具来帮助我们评估算力要求。
再回到硬件的算力，苹果的 A15 处理器仅 NPU 就具备 15.8 TOPS 的算力，而高通骁龙 855 处理器 CPU+GPU+DSP 叠加的 AI 算力仅有 7 TOPS，可见，要想算法模型可以流畅的在移动端运行，还需要对算法模型进行一些处理，让算法模型可以适配 Android 或 iOS 的硬件加速能力。
转换算法模型到移动端有很多种方法，大体上可以分为框架提供的转换功能和第三方转换工具两种。框架提供的转换功能优点是比较简单，可以直接通过 API 将模型进行转换，缺点是 Runtime 需要依赖于框架提供，而这些 Runtime 之间往往是不兼容的。第三方转换工具相对麻烦一点，为了能够转换不同的工具对模型输入有特殊的要求，输出方面也会有一定的限制，我们先看一下框架提供的转换功能。
全面的介绍所有框架的转换方法超出了本文的范围，我仅以 TensorFlow 提供的转换功能为例，介绍通常情况下转换模型的方法，这些方法在不同框架下大同小异，只需要查找文档中相似功能的 API 即可。
在 TensorFlow 中对模型进行转换主要分两种情况，一种是在调试、训练模型的过程中进行转换，另一种是在框架保存模型后进行转换。这两种情况的区别在于，在调试、训练模型过程中转换更简单和直观，有算子的兼容性等问题可以随时调整，保存后的模型文件转换则不能随时调整模型定义和算子等，但是，保存后的模型文件转换更容易做好工程链路，因为文件作为输入可以把模型训练和模型转换解耦开。我的建议是在端智能项目初期先用框架能力在调试、模型训练过程中转换，然后部署到移动设备上进行测试。
通过上面的代码，借助 tf.keras 的 API 定义了一个简单的模型，再借助 tf.lite.TFLiteConverter 将 Keras 模型转换成 TFLite 模型，最终保存成 .tflite 后缀的模型文件，我们的模型准备工作就初步完成了。
如果要把算法模型文件转换成 TFLite 用 TensorFlow 提供的命令行工具即可。
第一个参数是原模型文件的输入路径，第二个参数是转换后 TFLite 模型的输出路径。
由于 TensorFlow 框架在转换过程中做了一定的优化，因此，模型的 TOPs 也从 2.8 MFLOPs 降低到 1.8 MFLOPs。这里对 TFLite 模型算力评估使用了开源工具 tflite_flops ，按照下面的示例安装和使用即可。
本质上端智能用到的模型是从台式机、服务器上训练的模型转换过来的，这种转换在不同的移动端和边缘设备上是不同的，这就造成了转换的复杂性。
ONNX 作为开放神经网络交换标准，对不同框架和不同的移动端、边缘设备 Runtime 进行了标准化，从而降低了模型在不同框架和不同 Runtime 之间转换的成本，现在主流的框架和设备都支持 ONNX。
上面是 TensorFlow 命令行工具转换模型到 ONNX 格式的示例，当模型被转换成 .onnx 格式后，不仅可以在 ONNX Runtime 上直接运行，还可以方便的导入到不同设备的 Runtime 中运行，主流的 Android 和 iOS 设备的 Runtime 都支持对 ONNX 格式模型的导入，比如像下面这样导入 iOS 。
还有一种比较重要的方式就是编译，你可能会问：机器学习不就在框架的 Runtime 上解释执行么？要个编译过程有什么用？其实，这里大有学问，撇开兼容性针对不同 Runtime 的基础要求不谈，仅针对不同硬件的加速能力，都有一大堆工作需要软件工程师耗费巨量的精力。拿我当初在 UC 国际浏览器上做超分辨率的项目举例，为了能够让网络条件比较差的印度用户可以享受更高清晰度的图片和视频，我们训练了一个超分辨率算法模型，在中低端机型上经过优化的模型算法，基于 ARM 的 NEON 指令集加速，可以做到每秒 24 帧的速度把 240p 的视频实时转换成 720p。但是，这个项目最大的挑战不是算法模型，你在 Github 上一搜会有很多超分辨率 SOTA 的模型可供选择，压缩剪枝、量化、知识蒸馏、降低精度等优化方法也能方便的找到资料应用一遍，但这些优化并不改变神经网络对算力的基本要求，再加上对视频解码、内存和数据传输等性能消耗，想要在中低端机型上把视频超清化算法模型跑起来并非易事，近半年的时间硬生生把黄振这个算法工程师逼成了 Android 底层研发工程师，基于 NEON 指令集和 Android 开放的底层优化能力，一点点压榨机能、压缩算力消耗才把性能从每秒 3~5 帧提升到 24 帧。
基于编译的模型优化
随着机器学习技术生态的发展，为了降低优化神经网络的复杂度和成本，机器学习领域有所涉猎的大厂纷纷提出了自己的解决方案，PlaidML、TVM、JAX、MLIR等。Google 的 XLA (Accelerated Linear Algebra) 打破了图优化和算子优化分层优化的思路，XLA 分为两部分 HLO（High Level Optimizer） IR用来做后端无关优化，这里面既包含了神经网络底层计算图相关优化，也包含公共子表达式消除、强度缩减等传统优化技术。好处是可以尽可能地利用 llvm 编译架构面向多后端的优势，将设备相关代码生成和优化交给 llvm 来做。类似地，Facebook 的 Glow 也包含了这两层优化，glow 更侧重于多后端以及新型芯片，XLA 和 Glow 都侧重与 llvm 的结合充分利用既有优化手段。
JAX 的定位科学计算（Scientific Computing）和函数转换（Function Transformations）的交叉融合、训练深度学习模型外，同时具备：
- 即时编译（Just-in-Time Compilation）
- 自动并行化（Automatic Parallelization）
- 自动向量化（Automatic Vectorization）
- 自动微分（Automatic Differentiation）
下面用一个实际的例子，来学习一下如何借助 JAX 准备算法模型，首先是安装依赖。
然后打开 Jupyter Notebook 开始实验，先引入必要的 Python 包。
准备训练样本和验证集相关数据，这里还是以 mnist 为例。
接着，用刚才引入的 JAX 包提供的 API 定义模型。
这里和之前用 Keras 定义模型略有不同，但是细心的你一定发现了很多似曾相识的部分：Dense、Relu、Softmax，这些神经网络层和优化器定义和传统的机器学习架构提供的能力类似，就是用来定义和训练模型的。训练的具体步骤这里不赘述了，如果决定使用 JAX 可以去 Tensorflow 官网找到完整示例，接下来看一下模型的转换和保存。
依旧使用 TFLiteConverter 转换，把 jax 模型转成 TFLite 格式和之前介绍的模型转换并无不同，因此，也可以把 jax 模型定义和训练的过程用 Tensorflow.keras 进行替换，使用 jax 的好处是模型的训练会在 jax 的 JIT 环境中充分发挥 TPU、GPU、NPU 等硬件的加速能力。（具体的支持情况可以参考 jax 文档）
截止 2022 年初 JAX 仍然是一个实验性框架，而且对模型的编译和加速多集中在定义和训练模型环节，一旦模型转换成 TFLite 部署到手机上，依赖的是 TFLite 的 Runtime 而非 JAX JIT 这点必须有明确认知，避免错误的认为用 JAX 在端智能时会实现性能加速。这里介绍 JAX 的原因在于，如果你持续深入对机器学习进行实践，像我们在中低端 Android 手机做超分辨率一样，在性能优化上遇到诸多瓶颈，而 JAX 可以赋予你绕过 Keras 高级 API 进一步优化神经网络的偏底层能力，又不会太底层而造成复杂度激增。
另一个优化算法模型的工具就是 TVM，这个由 Apache 软件基金会赞助的开源机器学习编译框架，相较于 JAX 更注重面向 Runtime 的优化，尤其是对移动设备的支持能力更好且更易于扩展，除了传统的移动处理器和神经网络加速器，像 ZYNQ 这种 FPGA 的神经网络加速硬件都是支持的，甚至允许你自己用 DSP、CPLD 等定义自己的加速硬件并添加 TVM 的编译支持，TVM 会针对特定硬件自动优化模型。
虽然 TensorFlow.js 和 ONNX.js 将机器学习引入浏览器，在本书开始的部分介绍过基于 Tensorflow.js 的示例，但 Web 版本和原生版本之间在性能上仍然存在不小的差距。原因之一是缺乏对 Web 上 GPU 的标准和高性能访问，WebGL 缺少高性能深度学习所必需的计算着色器和通用存储缓冲区等重要功能。
WebGPU 是即将到来的下一代网络图形标准，它有可能极大地改变这种情况。与 Vulkan 和 Metal 等最新一代图形 API 一样，WebGPU 提供一流的计算着色器支持。为了探索在浏览器中使用 WebGPU 进行机器学习部署的潜力，TVM 以针对 WASM（用于计算启动参数和调用设备启动的主机代码）和 WebGPU（用于设备执行），在 Web 上部署机器学习应用的同时，仍然提供接近 GPU 上的原生性能。
图 4-6 WebGPU、Metal、OpenCLI 性能对比（摘自 TVM 官网）
比较通过 TVM 的 WebGPU 后端和使用原生 GPU 运行时（Metal 和 OpenCL）的原生执行完整计算图，在 MobileNet 模型上，我们可以发现 WebGPU 接近 Metal 的性能，假设 Chrome WebGPU 的运行时针对 Metal 而不是 MacOS 上的 OpenCL，你可以假设针对 GPU 的神经网络加速几乎没有性能损失。
此基准测试不包括 CPU 到 GPU 的数据复制成本，仅对 GPU 执行进行基准测试。目前从 CPU 到 GPU 的数据复制仍然需要 25% 的执行时间，但是，这些成本可以通过连续执行设置中的双缓冲等方法进一步摊销。
尝试 WebGPU 是为深度神经网络（矩阵乘法和卷积）中的原始运算符编写着色器，然后直接优化它们的性能，这是 TensorFlow.js 等现有框架使用的传统工作流程。TVM 则不同，它采用基于编译的方法来提供接近 Native 的性能。TVM 能自动从 TensorFlow、Keras、PyTorch、MXNet 和 ONNX 等高级框架中提取模型，并使用机器学习驱动（借助机器学习的能力让 AI 来决策如何优化）的方法自动生成更为优化的 Native 代码。
基于编译的方法的一个重要优点是底层的复用，通过复用底层来优化目标平台（如 CUDA、Metal 和 OpenCL）的 GPU 内核，TVM 能够毫不费力地针对 Web 进行优化。如果 WebGPU API 到Native API 的映射是有效的，只需要很少的工作 TVM 就可以帮我们在 Web 上达到类似的性能。更重要的是，AutoTVM 基础架构允许为特定模型专门化计算着色器，从而为我们感兴趣的特定模型生成最佳计算着色器。（这种自定义能力和前面提到的针对自定义硬件加速是一脉相承的）
为了构建一个针对 Web 且基于 WASM 和 WebGPU 的端智能程序，需要以下部分：
- 用于计算着色器的 SPIR-V 生成器。
- 主机程序的 WASM 生成器。
- 加载和执行生成的程序的运行时。
幸运的是，TVM 已经为 Vulkan 提供了 SPIR-V Target，并使用 LLVM 生成主机代码，所以我们可以重新利用这两者来生成设备和主机程序。
TVM 具有最低限度的基于 C++ 的运行时，构建了一个最小的 Web 运行时库并将其与生成的着色器和主机驱动代码链接，生成单个 WASM 文件。通过在 TVM 的 JS 运行时中构建 WebGPU 运行时，并在调用 GPU 代码时从 WASM 模块回调这些函数。使用 TVM 运行时系统中的 PackedFunc 机制，如图 4-7 将 JavaScript 闭包传递给 WASM 接接暴露
高级运行时原语。
在 WASI 的和 WASM 的帮助下，TVM 能够很方便的让前端把一些流行的 AI 算法模型编译到 Web base 的环境，并且有良好的性能加持。
如图 4-8 所示 TVM 的架构更具灵活性，尤其是对端智能领域不同的操作系统、硬件设备加速能力，开发人员不必再手动进行适配和优化，极大提升了端智能工程的效能。下面用一个具体示例演示使用 TVM 的方法，便于您评估该方法是否契合您的技术工程体系。
在开始之前，我们先从 Github 上下载一个 ONNX 的模型。
模型地址有可能变化，因此，如果下载失败可以在 https://github.com/onnx/models/ 里重新找到并下载模型。
接着，需要安装一些必要的软件，这里以 MacOS 为例介绍安装过程，其它系统可以在 TVM 的官网 tvm.apache.org 找到.
安装完毕后下载源码并生成配置文件。
这里需要注意的是对编译选项的修改，如果想要在本地测试要把 set(USE_LLVM ON) 这个默认 OFF 的开关改成 ON。如果想启用 CUDA 后端，要为 TVM 构建其他后端和库（OpenCL、AMD 生态 RCOM、APPLE 生态METAL、VULKAN ......）对开关 set(USE_CUDA ON) 进行使能配置即可。
开始编译时注意 -j4 这个编译选项，根据自己的线程和 CPU 核数来合理设置，然后等待编译结果即可，我这边是一次性就编译完成，用的 MacOS 12.2.1 版本系统。接下来，进行 Python 库的安装。
需要注意的是 TVM 的优化分为三层，第一层是直接 TVM 编译后的模型会有一些编译层面的优化，第二层是利用 autotuner 这个 TVM 提供的自动优化利器来针对不同的 Target 如：CPU、GPU、NPU、TPU 等生成优化的代码，第三层是利用 TVM 的分析工具和自定义调整能力来进一步优化模型性能。通常在端智能领域能掌握第二层优化能力，就能够应付大多数场景了，还是要把精力平衡到输入处理、CPU、内存、I/O 等优化中去。
在开始之前还需要准备两个 python 脚本，分别用于生成模型预测的输入和解析模型预测结果，便于对 TVM 的编译和优化结果进行校验。
用 TVM 运行我们之前下载的图像分类模型进行预测，还需要对预测结果进行解析
下面就进入前述的第一层，用 TVM 的工具进行模型的编译。
让我们看一下模块中创建的文件：tvmc compile
您将看到列出的三个文件。
mod.so是模型，表示为 C++ 库，可由 TVM 运行时加载。
mod.json是 TVM 计算图的文本表示。
mod.params是一个包含预训练模型参数的文件。
有了模型和输入数据，我们现在可以运行 TVMC 进行预测：
回想一下，.tar模型文件包括一个 C++ 库、对 Relay 模型的计算图文本表述以及模型的参数。TVMC 包括 TVM 运行时，它可以加载模型并对输入进行预测。运行上述命令时，TVMC 会输出一个新文件 ，predictions.npz其中包含 NumPy 格式的模型输出张量。
在此示例中，我们在用于编译的同一台机器上运行模型。在某些情况下，我们可能希望通过 RPC Tracker 远程运行它。要了解有关这些选项的更多信息，请查看。
下面用后处理预测结果，用图 4-9 作为输入检查一下模型的分类预测准确性。
TVMC 将针对模型的参数空间执行搜索，尝试不同的配置并在指定的 target 上生成最快的参数。虽然这是基于 CPU 和模型的引导式搜索，但仍可能需要很长时间才能完成搜索。此搜索的输出将保存到 resnet50-v2-7-autotuner_records.json 文件中，稍后将用于编译优化模型。
如果为标志指定更具体的优化 Target，将得到更好的结果--target llvm -mcpu=skylake。前面的 -mcpu 参数在 Intel 处理器上可以使用，LLVM 指定 CPU 架构的编译进行调优。在 MacOS 上可以用命令 sysctl machdep.cpu 来查看处理器的详细信息，再根据信息去处理器官网来查询架构代号。
经过漫长等待后应该能得到如下信息，并得到 .json 文件所包含的调整数据。
现在已经收集了模型的调整数据，我们可以使用优化的运算符重新编译模型以加快计算速度。
验证优化模型是否运行并产生相同的结果：
验证预测是否和之前的结果相同：
TVMC 提供模型之间的基本性能基准测试工具，并且 TVMC 会报告模型运行时间。我们可以大致了解调优对模型性能的提升程度。例如，在我的 intel 处理器的本地测试，调整后的模型比未调整的模型运行速度快 30% 以上。
这些数据可能和你本地的测试有一定的偏差，TVM 官网示例中优化结果大概在 47% 左右的提升，因此，根据 target 的不同和模型不同，优化的结果可能会有差异，但综合来看提升的效果还是非常明显的，所以我们不必花费精力在第三层上做更深入的优化也能达到够用的状态。同时，必须要注意在编译和构建 TVM 的时候编译选项会直接影响可使用的 target 选项，我在这里也踩了坑重新编译安装了一遍，所以在最初编译安装的时候需要对未来移动端 target 有一个预估。
