首先讲一下我自己的整个经过,不想看到可以直接跳到标题处。
我先将在电脑上训练的pytorch模型转换成了ONNX格式的模型,再经过开源的工具转换成了c语言代码,再编译进了项目里,最后部署到了RV1126里,如下是效果图,
这是训练了一个ai贪吃蛇,能够保证在吃的足够多的情况下自主的吃豆子,接下来我用同样的方式将该模型部署到esp32s3上,结果发现模型每运行一次需要3s,这个运行速度是完全不行的,后来了解到要对模型优化才能够部署到mcu上,有多种方式优化,我最终选择的是对模型进行量化,量化就是将浮点型数据量化为整形数据,比如32位浮点型量化位8位的整形,从存储空间上来说就缩小了4倍,再说运行速度,MCU处理8位整型也是要比32位浮点型要快的。然后我先对转换后的ONNX格式的模型进行量化,结果发现量化后无法通过onnx转c语言的工具正确识别,然后我又先对pytorch格式的模型先量化再转换成ONNX格式,发现还是不行,由于我是刚接触ai,对这些不太熟悉,无法去修改,找了其它的开源项目也没有办法,3秒的时间又无法接受,可以想象一下贪吃蛇每3s移动一下的画面,所有就先搁置了。
后面过了一段时间发现了tensorflow lite的项目,官方有pytorch转tflite的工具,而且乐鑫也有对tflite支持的库,并且用到了自身的esp-nn库,硬件加速ai,我就试了一下,最后将tflite格式部署到esp32s3上,运行一次的速度要1s(不清楚有没有实际启用esp-nn),这个速度比之前onnx快了2s,但还是有点慢,接下来我就对tflite量化,问题就出现在这里了,tflite格式的模型不支持量化,只能在转换为tflite前量化,我就先对pytorch量化,在转换tflite的过程中报错,看不懂报错说明,而且我也认为量化要用自己的架构会更好,只能换另一种方式,先将pytorch转换成tensorflow模型再量化,结果在这里卡了好久,在网上找的工具转换时各种报错,最后多次尝试后终于解决了,总结下就是pytorch先转换为ONNX,再转换为tensorflow,之后再量化,在量化途中就直接转换为tflite格式,最终运行一次的速度为100ms,搞定了,这是目前已经可以接受的速度。
下面就详细的讲一下过程。
<