本文介绍 2022_06_25_10_探索一条模型部署新思路：如何手术级精准操作你的模型

2022_06_25_10_探索一条模型部署新思路：如何手术级精准操作你的模型

前段时间我设计了一个推理IR，包括了一些常用的算子集合，算是把工业界部署的一些逻辑、模型的构建方式方法走了一遍。感兴趣的朋友可以去看我之前的一个完整的系列文章。

但是只有IR并没有什么卵用，于是我基于这一套IR尝试构建一个推理引擎，事实上这件事情是可行的，我通过简单修改了一下netron的代码，可以实现这套IR的可视化，比如可视化一些简单的模型，例如几层FC没有任何问题：

现在已经可以支持残差模型的结构了：

这套IR的个人认为最大的优点在于，非常精简，而且很好维护，这比ONNX不知道简单多少倍。当然最核心的并不是这个IR的设计，而是实现了一套直接通过torchscript到WNNX的流程。这也是设计这套IR的初衷，并不是为了造轮子，而是的确没有好用的轮子，随着pytorch的生态不断的扩展，pytorch本身的模型表征格式在不断的完善，那么越贴近训练框架肯定是越好的，这也是我一开始设计就是直接从torchscript到WNNX格式，砍掉中间商，实现端到端对接，让你的模型训练什么样，部署就什么样，而不至于ONNX一倒出来，就面目全非。

然而，这件事情并不容易，我做了几个星期，虽然能把一些简单的CV模型倒出，例如VGG，Mobilenetv2，Mobilenetv3等都没有问题，但要变成一个真正可用的框架却还有一大段路要走。比如，你需要一个推理引擎对其进行推理，而这个引擎要处理的问题就更加复杂了：每个算子要对齐结果，要处理很复杂的网络图链接，要实现每个算子的算法并且至少要支持X86 和 ARM平台架构，如果要做的完善一点，还要支撑起诸如FP16量化，int8量化等常规操作。这里面 每一点都可能要耗费我一年的时间。

尽管如此，也不是没有好消息，其中之一就是我们不需要兼容太多的框架，不需要像MNN NCNN等那么大而全，根据当初设计的初衷，很多算子的实现只需要按照pytorch的实现对齐就可以了，这也大大减少了工作量。

由于上述的种种原因和复杂性，直接就导致了我的WNNX系列文章的最终章 - 实现YOLOX的推理迟迟没有推出，因为要对齐网络的推理结果实在是太难了，但是挖下的坑一定得填完，大家一定要对此抱有期望和信心，如果我没有搞出来，那你可以取关我，如果我搞出来了，请帮我拉十个粉丝。

虽然跑YOLO这个目标还未达成，但我 发现这个IR还有一个额外的功效，这使得我开始思考一条模型部署的新思路，使用这个额外技能，我们似乎随心所欲的操控我们的模型，彻底摆脱工具链的依赖，完成 任何我们想实现的操作，就跟你调用函数一样简单。

思考

相信很多同学部署模型肯定上来就是ONNX导出一波，额，不对，貌似现在是直接MMDeploy上来搞一波，拿来就用，甚至都不用思考。但是你有没有想过，其实要达到你的目标，路径并不是一条，如果你选择了容易的路，那可以会错过很多迷人的风景。

这种方法简单易行，但是缺乏灵活性。举个例子：

你有没有想过这种操作其实是非常的暴殄天物？对于GPU倒还好，但是对于计算资源非常有限的场合，这完完全全没有必要。那为什么现在很多库号称可以做optimization，可以做图融合，却不能解决这样的问题呢？如果你深入去看那些optimize的代码你就会发现，这些融合的操作是非常有限的，甚至都是基于已知的模板优化，换一个模型，换一个稍微复杂点的操作，就支持不了了。

所以这就是为什么得出这个结论：ONNX的方式，适合不需要思考的场景，上来直接撸起袖子开始干的场合。 ，但是对于一些更小的模型，需要外科手术级进行操作的模型，这种方式就不太适用了。

新思路

我在做这套IR的时候，猛然发现它的功效其实还可以做很多事情，不仅仅是学习，甚至可以完成一些工业落地的东西。介绍一种模型部署新思路：我们有没有一种可能，直接用推理框架来构建模型？答案是完全可以的。直接上代码：

熟悉TensorRT的同学可以一看就知道了，我们基于WNNX实现了一套NNBuilder搭建的API！！这跟用TensorRT API构建网络的方式是不是很像？但是要简单的很多！而且这个serialize的模型，是可以直接被可视化的：

可以看到，生成的模型结果，跟你从onnx导出，是一致的，并且推理结果也完全一样。虽然看起来很容易，其实这背后要处理的逻辑还是比较复杂的，例如要构建图，要处理多输入多输出等复杂场景，甚至对于我们这个WNNX框架，你还得做Shape inference，否则推理性能得不到保证。总之，TensorRT的 networkBuilder能做的，我们都能做。

尽管一些主流推理框架也有这个功能，例如MNN，ncnn，但API的易用性、方式的简单性，都不如我们自己的做的。虽然TensorRT门槛高较难用，但在API构建网络方面非常值得学习借鉴，现在，你也可以用NNbuilder来搭建一个基于CPU的模型了。

手术级操作

上面的新思路，虽然有些框架有，但是从易用性层面来讲，能做到 可行、可用、易于理解的不多。那么基于上面的API，我们怎么才能够实现手术级操作模型呢？

其实还有两个核心操作：

直接操作模型每一个layer的权重；
动态控制输入输出逻辑，动态控制流图。

操作起来也非常简单，例如上面的全链接模型，我们直接把权重保存下来不就行了。然后直接：

auto w0 = wnn::from_numpy_bin("a.bin", {12, 4});

WNNX提供了 numpy 接口，你可以直接从numpy保存的bin文件读取权重。这不比 1+1 还简单？

与此同时，你甚至可以直接一键量化：

if (fp16_mode) {
    mBuilder->set_fp16_enable();
    mBuilder->serialize_to_file("created_fc_fp16.wnnx");
  } else {
    mBuilder->serialize_to_file("created_fc.wnnx");
  }

目前支持的量化操作包括但不限于：Convolution，FC，以及一些自定义算子。当然了，int8就稍微复杂点了，有fp16也够用。

总结

本篇给大家介绍了一种模型部署的新思路，通过直接操作权重和网络流图的方式进行模型搭建，这一些操作的可能性是基于我们有一套简单易用的网络搭建API，借助这个强大的工具来实现复杂的逻辑操作，实现控制一些更复杂的小模型也成为可能。

而这种手术级的操作能力，也让你可以训练模型用一套，部署的时候那些层我想要，那些不想要，哪些权重我可以放，哪些不需要放，哪些可以融合，哪些可以拿掉，掌控权全部在你。

这也就，实现了从底层搬砖码农到外科手术医师的华丽转身。