本系列文章讲述的是如何从0构建一套自己的推理引擎,真正的硬核科技,如果你还没有阅读我之前写的三个篇章,可以关注我,从我的知乎主页或者我的公众号上进去阅览。
今天我们开始着手下一个工作:验证我们初版推理引擎的正确性。
我这里对一个非常简单的线性模型,也就是全链接层的模型进行了验证,pytorch原始模型的结构如下:
这个模型非常简单,正好用来验证我们的推理引擎的内存管理、层调度、模型解析方面的正确性。首先我们先将这个pt转到 **wnnx**
:
./build/src/townnx -m test_simple_fc1.pt -s 1x864
这个生成的wnnx模型也可以用 netron
简单的可视化:
可以说超级无敌的简单,连彩色都没有。。。
我们可以看到输出:
大概我们可以看到pt模型被保存为了 wnnx
的格式,并且可以通过输出直接看到这个模型的具体信息,例如每一层的名字,类型,甚至是权重信息,输入输出等。
**这个就很方便了,顺便说一下,通过 ****townnx**
这个工具,其实也是可以直接丢进去一个 wnnx模型的,然后它会吧wnnx以文本的形式打印出来,可以说非常的方便。
接下来就是推理了:
可以看到我们的推理引擎在逐步的构造每一层,并且从wnnx文件中去解析每一层网络结构以及权重。
接下来见证奇迹的时刻到了,看看模型的输出结果:
我们直接查看最开始的40个元素,上面是wnn的推理结果,再来看看pytorch的:
和pytorch的推理结果一致!
成功了!来总结一下我们成功的完成了哪些事情:
- 我们创建了一套自己的IR,并且可以直接从torchscript转到它,在这个过程中就完成了很多图简化的过程;
- 我们的WNNX IR具有轻量化,单个文件保存所有信息,可加密等优点,并且我们使用netron改造,实现了对wnnx模型可视化的支持!
- 最重要的是:我们完成了一套非常轻量级的WNNX推理引擎,它可以读如wnnx格式模型文件,并完成推理,结果与pytorch结果一致!
尽管如此,但看起来我们只是造了一个新的轮子,接下来我们要思考一下,我们怎么将其变成一个有特色的推理引擎呢?
我的回答是:
- 简单化;
- 轻量化;
- 以及更简单的量化;
- 专注transformer推理优化。
**简称四化。**为什么说这四个方向可以让我们打造的这个推理引擎变得更有特色呢?举个很简单的例子,我们的框架已经实现了 Conv2d, FC (InnerProduct), Pool2d, Laynorm各种激活函数等层, 以及基本的图调度,Tensor管理,输入输出结构,用这些东西足够推理一个VGG16了,而我们库到目前为止:
还没有超过1M!谁能抵抗一个这么小又可以推理VGG16这种比较大模型的推理引擎呢?你封装到APP里面大小也才跟一个txt差不多!但是它却可以Power整个AI模块!
最后一点,也显而易见,现在的推理框架对于transformer的优化都太不友好了,这里面有很多工作是可以做的,不仅仅是运算方式的加速,也包括transformer模型本身一些attention的简化设计。最近有个工作:
EfficientFormer就把这个议题提上了议程。
好了,今天的牛逼先吹到这儿。
接下来要做的事情是让前推支持多个输入。相信我,这件事情没有想象的那么简单,甚至有些许的复杂。因为要考虑的情况着实是太多。下一期我们继续来深入解决一下这个问题。