浅谈机器学习模型推理性能优化

在机器学习领域，清晰明了的数据预处理和表现优异的模型往往是数据科学家关注的重点，而实际生产中如何让模型落地、工程化也同样值得关注，工程化机器学习模型避不开的一个难点就是模型的推理（Inference / Serving）性能优化。

可能许多数据科学家都对模型的推理性能比较陌生，我举几个对推理性能有强要求的场景例子：

1. 在公共安全领域中，视频监控中实时的人脸识别需要有实时的展示能力方便执法人员快速定位跟踪人员。

2. 在互联网应用领域中，电商网站、内容应用实时的个性化推荐要求能够快速响应，推荐的卡顿感将直接影响购物或者内容获取的体验。

3. 在银行领域中，电子支付中异常交易的实时识别也至关重要，任何异常的交易需要被快速识别并拦截，而正常的交易则不能被影响。

4. 在金融领域中，量化模型毫秒级的交易判断输出能帮助华尔街的交易员们套取巨额利润。

从上面的例子不难发现，其实在不同的领域的场景下，推理的性能都是模型表现之外最关注的点，在某些极端的场景，数据科学家和机器学习工程师甚至愿意牺牲一部分的模型表现来换取更高的推理性能。

计算图优化

“提高硬件性能是优化的最后一步，而不应该是第一步。”

上面这句话是我们项目 PO 说的话，其在 Spark 性能优化上有非常丰富的经验，我非常赞同这种论点。据我观察，在遇到算法模型的训练和推理性能瓶颈的时候，大部分机器学习工程师都希望能获得更高的硬件性能来突破瓶颈，却忽略了计算逻辑本身的优化。更高性能的硬件为模型推理带来的性能提升并不是线性的，而花费的硬件成本却是指数级上升的，所以一定要记得，不到万不得已，千万不要指望硬件带来的性能提升。

基本上数据处理和算法模型都可以被抽象为计算图，而计算逻辑的优化往往在领域内被称为图优化（这里的图优化并不是指图模型的表现优化哦 :D）。

每个计算图中都包含许多计算节，图优化的目标很简单，就是简化计算图中计算节点的计算量。常用的方式分为以下几种：

1. 减少节点的数量

2. 用高效替换低效的节点

3. 用高效子图替换低效子图

4. 用并行化分支代替单分支

减少节点的数量

在构造机器学习模型的时候，我们往往会无意中对数据做了多余或者反复的操作，这类操作就像写工程代码中的 code smell 一样，在模型构造完成之后一定要对这种操作多加注意。拿矩阵的转置（transpose）做例子（实际上多余、反复转置是非常常见的）：

def func(a, b):  a_T = a.transpose(1, 0)  c = a_T + b  return c.transpose(1, 0)  def func_better(a, b):  return a + b.transpose(1, 0)

In [10]: a = torch.randn(10000, 50000)
In [11]: b = torch.randn(50000, 10000)
In [12]: %timeit func(a, b)9.67 s ± 827 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
In [13]: %timeit func_better(a, b)8.44 s ± 233 ms per loop (mean ± std. dev. of 7 runs, 1 loop each）

用高效的节点替换低效的节点

相同的计算节点往往有多种实现方式，而这些实现方式中往往都有各自的优劣势，有一些是牺牲了空间换取时间，而有些是牺牲了时间换取了空间，如果是考虑推理的响应性能，那么我们往往会用时间有优势的实现方式来替代时间没有优势的节点。

def func(a):  result = 0  for i in range(a.shape[0]):    result += a[i][i]  return result
def func_better(a):  eye = torch.eye(a.shape[0], dtype=torch.bool)  return a[eye].sum()

In [17]: a = torch.randn(1000, 1000)
In [23]: %timeit func(a)6.75 ms ± 72 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)
In [24]: %timeit func_better(a)2.48 ms ± 40.1 µs per loop (mean ± std. dev. of 7 runs, 100 loops each

用高效子图替换低效子图

# Example:# a_original -> [[1,2,3], [4,5,6], [7,8]]# a -> [1, 2, 3, 4, 5, 6, 7, 8]# row_idx -> [0, 0, 0, 1, 1, 1, 2, 2]# row_lens -> [3, 3, 2]
def func(a, row_idxs):  result = []  for i in range(row_idxs.max() + 1):    result.append(a[row_idxs==i].sum())  return torch.tensor(result, dtype=torch.float32)
def func_better(a, row_lens):  result = []  for each in a.split(row_lens.tolist()):    result.append(each.sum())  return torch.tensor(result, dtype=torch.float32)

In [66]: a = torch.randn(100000)
In [67]: row_idxs = torch.tensor([0] * 30000 + [1] * 50000 + [2] * 20000, dtype=torch.int32)
In [68]: row_lens = torch.tensor([30000, 50000, 20000], dtype=torch.int32)
In [69]: %timeit func(a, row_idxs)1.75 ms ± 9.35 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
In [70]: %timeit func_better(a, row_lens)43.4 µs ± 412 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each

用并行化分支代替单分支

在许多图计算框架里面，并行加速分为两种，一种是算子内部的并行化（intra）、一种是图分支的并行化（inter）；例如在 ONNX 中，一个 for 循环算子是无法得到并行优化的，因为其维护了一个状态变量 i，而往往我们并不会使用到这个 i，我们只是想让某个计算逻辑执行 n 遍。这个时候就可以将这个 for 循环算子，拆分成 n 个计算分支，这样在使用图分支并行化计算的时候，就可以充分利用硬件资源提高计算效率了。

Backend 优化

计算图优化是第一步也是最重要的一步，那么在计算图优化和硬件优化之间，难道就没有其他优化方式了吗？答案是有的，这个方式就是计算图引擎 Backend 的优化。计算图只是一个计算逻辑的抽象表示，而真正执行计算图的引擎也会有不同的实现，而每种实现往往带来的都是不同的计算效率。比如在 PyTorch 模型的 Inference 的可选项里面，有以下几种计算图引擎后端可供选择：

1. 原生 PyTorch API

2. TorchScript Python API

3. LibTorch

4. ONNX Runtime

原生 PyTorch API

原生的 PyTorch API 其实不用过多的描述，就是执行 PyTorch 模型中的 forward 函数，直接得到推理结果。可以说这个 API 是最简单、最原生的方式，可以作为推理性能表现的的一个 BaseLine。这个 API 推理过程和训练保持一致，可以保障结果的正确性，可以作为其他 backend 正确性检验的一个对照。

TorchScript Python API

TorchScript 是一种 PyTorch 模型的表示格式，相当于 PyTorch Python API 的子集构建的子语言，其能够被 TorchScript 编译器实时编译成 C++的模型代码并执行。这种格式有三个主要的设计初衷：

1. 构建一种跨环境序列化模型的方式

2. 基于 Torch 基本算子，并可扩展的算子集

3. 可以在 C++程序中实时执行

通过 torch.jit.script 的 API，可以将一个 Python 模型转换为 TorchScript 模型，并通过 torch.jit.save 保存为.pt 格式的 TorchScript 模型。注意：转换成 TorchScript 的 PyTorch 模型要求用一定的规范编写，可以参考官方的文档：https://pytorch.org/tutorials/beginner/IntrotoTorchScript_tutorial.html?highlight=torchscript

对于 TorchScript，PyTorch 是有 Python 的 API 支持的，通过 torch.jit.load 可以读取一个.pt 的模型文件，并执行 forward 函数即可进行推理服务。然而，实际测试发现，这种推理性能与 PyTorch 原生的 API 的性能是较为接近的，仅仅是在稳定性有小幅度的领先。

LibTorch

既然 TorchScript Python 的 API 那么弱，那我们就来试试 C++的 API 吧！TorchScript 的 C++API 是 LibTorch， 使用 LibTorch 编译后的推理性能无论是速度和稳定性能有显著的提升。官方提供了一个简单的例子进行参考https://pytorch.org/tutorials/advanced/cpp_export.html。

有个小疑问：实际在客户现场的 Linux 服务器上，LibTorch 的表现稳定性相当差，而在我自己的 MacBook 上是很稳定的，不清楚是什么原因。我怀疑和编译器或者基础库有关

ONNX Runtime

有人会问了，既然那么好，为啥不能成为唯一的选项？答案是 LibTorch 是 C++的库，对编译和环境的依赖比较严重，并且对 C++编程水平的要求会比较高。所以接下来我要讲讲我们最后选用的方案———ONNX Runtime。

ONNX(Open Neural Network Exchange)是一种针对机器学习所设计的开放式的文件格式，用于存储训练好的模型。它使得不同的人工智能框架（如 Pytorch, MXNet）可以采用相同格式存储模型数据并交互。ONNX 的规范及代码主要由微软，亚马逊 ，Facebook 和 IBM 等公司共同开发，以开放源代码的方式托管在 Github 上。目前官方支持加载 ONNX 模型并进行推理的深度学习框架有：Caffe2, PyTorch, MXNet，ML.NET，TensorRT 和 Microsoft CNTK，并且 TensorFlow 也非官方的支持 ONNX。——Wikipedia

针对这种通用的交换格式的模型，微软牵头发起了 ONNX Runtime 的项目 这个项目旨在直接运行 ONNX Runtime，相当于纯的模型推理 Backend，其设计的理念就是为了解决训练和推理的性能问题，并且支持各种硬件加速库加速（如：MKL、CUDA、TensorRT 等等）。除此之外，ONNX Runtime 还有 Python、C++、JAVA 等多种接口；甚至提供了直接用于 Serving 的程序，暴露了 HTTP2.0 和 GRPC 的接口，用起来非常方便。

PyTorch 提供了模型转换为 ONNX 模型的接口 torch.onnx.export，通过这个接口我们就可以将模型转换为 ONNX 模型在 Runtime 中进行推理了。通过测试，我们发现 ONNX Runtime 在推理的速度和稳定性上都是相当优秀的。略有些遗憾的是，PyTorch 中有些比较酷炫的算子 ONNX 并不支持，不过 ONNX 才刚刚兴起，相信之后一定会加入更多好用的算子的。

性能对比

• 机器：MacBook Pro (Retina, 15-inch, Mid 2015)

• CPU：2.2 GHz Intel Core i7

• 内存：16GB 1600 MHz DDR3

• device：CPU

• 模型：SqueezeNet

• 数据：(5, 3, 64, 64)

• Loop: 1024

原文链接：浅谈机器学习模型推理性能优化