PNNX 相关内容介绍

本文主要是根据 b 站关于 PNNX 的分享视频做的一个内容摘录，如果感兴趣的同学可以移步☞： PNNX介绍

本人还没人正式地使用过 ncnn 这个开源项目，但是视频中提到的 nihui 大佬的一些总结还有对于深度学习框架的观点看完以后还是大有受益的。也是引起了对 ncnn 以及最新的 PNNX 的一些兴趣，之后有机会可能会带来一些相关内容的博客。

因为 PNNX 是 ncnn 针对 PyTorch 框架所做的一个模拟转换工具，并且是脱离了采用 ONNX 这种中间表示的，所以一下的内容主要涉及到 PyTorch 和 ONNX 框架。

PyTorch 部署方式

PyTorch 有多重部署方式，而它们之间各有优缺点吧，这里主要形式如下：

1. PyTorch -> Libtorch
   即使用原生支持的 Libtorch 来完成部署，这样的好处是 Libtorch 对 PyTorch 的算子都基本支持，但是速度并没有厂商专门优化的库快。例如在 CPU 上会比 OpenVINO 慢，在 GPU 上会比 TensorRT 慢。并且 Libtorch 这个库比较大，对一些端侧并不是很友好。
2. PyTorch -> ONNX -> onnxruntime
   即 PyTorch 模型先导出到 ONNX，然后使用 onnxruntime 来运行，好处是 onnx 支持的框架多，onnxruntime 也可以支持一些 tensorRT 这样的后端，缺点就是 PyTorch 部分算子在 ONNX 中没有，以及 PyTorch 一个算子可能会导出 ONNX 很多的胶水算子，计算图变复杂，导致推理效率下降
3. PyTorch -> ONNX -> 第三方后端（主流方式）
   这里的第三方后端有 TensorRT，ncnn，OpenVINO，TNN 等，这样的优势是不同的后端在目标平台都有相对应的优化，可以获得最快的推理速度，但问题是第三方后端（库）对 onnx 算子的支持优先，会有运行不了的问题。
4. PyTorch-> TorchScript -> 第三方后端
   主要有 NV 的 TRtorch 项目，直接加载 PyTorch 的图进行转换。

Lower

Lower 是指将不支持的算子映射到一个或者多个算子上，这样用一些有限的基础算子，就可以拼接成很多新的算子功能了。

在 PyTorch 原生算子到 TorchScript 就有 lower 的过程，并且 TorchScript 到 ONNX 也有 lower 的过程，这两层 lower 就有可能会层层调用，得到一个非常庞大的计算图。算子粒度太细，并且得到之后的图不能再对应回原先算子了。

针对上面因为 lower 得到复杂计算图的情况，ONNX 社区 daquexian 有贡献一个 onnx-simplifier 工具可以用来简化图，将一些小算子映射回大算子，但是这一类的方式当前主要是做一些图优化并且基于 pattern 匹配，这也有几个缺点，首先是 ONNX 本身的表示不是很好的支持图优化，所以需要大量的 if else 逻辑来判断；第二点则是每次 PyTorch 或者 ONNX 版本不一致，同一个算子生成的计算图可能都会不同，pattern 会失效。下游框架开发者会很麻烦，每次升级不匹配都要再重新写一个图优化。第三点是高层算子输入的变化也会导致图的变化（例如同一个算子不同的参数或者属性组合映射成的图是不同的），这样也会增加图优化匹配的复杂性。

一些思考

为什么要使用 TorchScript 和 ONNX 来让原始的图变得更加复杂呢？一个好的 PyTorch 模型转换需要有哪些特点呢？
作者给出了一些思考点：

• 研究人员更喜欢高级别的 IR 表示
• 库和运行时也会喜欢高级别的 IR 进行图优化
• 人类可读可编辑的表示会更友好
• 更接近 PyTorch 原生态

PNNX

PyTorch Neural Network Exchange，PNNX，其在 PyTorch 在 ncnn 部署的流程大概如下：

即 Pytorch -> TorchScript -> PNNX -> ncnn， 并且 PNNX 还可以再转回 Pytorch python 代码，其为了综合上面总结的高层次，PNNX 的接口与 PyTorch 的 python api 一模一样，做到了完全一致的对应关系，这也是可以直接映射回 PyTorch python。

PNNX 延续了 ncnn 保存模型的形式，分成了 Param & Bin，并且可以纯文本来编辑模型，当前有一下几个特性：

• 提供一些 Module pattern 的匹配特性，可以用户指定多个小算子映射成一个大 op 的特性
• 支持 custom Op 的导出
• 对 QAT 支持性比较好（ONNX 只能支持部分）
• 支持静态 shape 和 动态 shape

其中有 GraphRewritter 类用来做 PNNX 转回原生框架的主要类。

关于 MLIR

在考虑全新中间表示的时候，nihui 大佬也考虑过 MLIR，首先 MLIR 接口还不是很稳定，并且作为 LLVM 的子项目，还需要编译整个 LLVM，然后 MLIR 是属于 Google 的 TF 系，TF 跟 PyTorch 完全是两个生态，所以在两个生态之间搭线的事不是很有趣，因此选择了一种更纯粹的直接接口映射。