菜单导航

我用深度学习做个视觉AI微型处理器!

作者: 精装之家 来源: 精装之家 发布时间: 2022年02月25日 09:35:15

讲多了算法,如何真正将算法应用到产品领域?本文将带你从0用深度学习打造一个视觉AI的微型处理器。文章含 完整代码,知识点相对独立,欢迎点赞收藏,跟着本文做完,你也可以做一个自己的嵌入式AI小产品!

属于AI人的硬核快乐

背景

随着硬件尤其是显卡性能升级,以及Pytorch,TensorFlow深度学习框架日趋完善,视觉AI算法在多个领域遍地开花,其中就包括嵌入式设备。这是一种微型处理器,它的关键单元就是内部小小的计算芯片。嵌入式设备和我们日常用的电脑相比体积小,只包含必要外设。一些针对特定任务的嵌入式设备往往不会运载我们常用的比如Windows、Linux系统,而是直接将代码烧录进去运行。

在嵌入式设备上尝试部署深度学习算法开始较早,1989年一家叫做ALVIVN的公司就将神经网络用在汽车上了。现如今,工程师们将其用在安防、机器人、自动驾驶等领域。因此,懂得如何设计、训练算法,又能将其部署到边缘硬件产品上,能帮我们实现许多产品的想法。

但是,视觉算法部署在产品中仍有许多难点,比如:(1)模型通常需要在CPU/GPU/NPU/FPGA等各种各样不同类型的平台上部署;(2)嵌入式算力/内存/存储空间都非常有限;跑在云端服务器上,需要实时联网又不很优雅;(3)模型训练时可能会使用不同的AI框架(比如Pytorch/TensorFlow等)、不同硬件(比如GPU、NPU),相互适配产生问题[1]。

因此笔者开始思考下列问题:

有什么亲民价格的芯片能处理部署视觉AI算法?

如何将深度学习算法部署到嵌入式设备上?

对第一个问题,在经过调研后,还真有这样的芯片,那就是嘉楠科技的K210芯片。一个芯片几十元,对应的开发板在某宝上两百多就可以买到。根据嘉楠官方的描述,K210具有双核 64-bit RISC-V RV64IMAFDC (RV64GC) CPU / 400MHz(可超频到600MHz),双精度 FPU,8MiB 64bit 片上 SRAM(6MiB通用SRAM+2MiB的AI专用SRAM)。关于这块芯片更详细的介绍可以参考[2] 。

市面上有许多搭载K210的开发板,笔者这里选了雅博一款功能较全的K210开发板,开始了嵌入式AI的折腾之路。

对于第二个问题,方法就多了,不同深度学习框架,不同硬件选型都决定着不同技术路线。基本路线可以为 深度学习平台训练 -> 模型剪枝、量化 -> 参数转换 ->转换为硬件平台上能运行的模型

对深度学习平台选型,笔者决定选用当下最流行的Pytorch平台。最后一步往往取决于这个硬件的生态,如果没有相关生态支持,可能需要手写C语言代码加载参数运行。调研发现,K210有一个深度网络优化平台 NNCASE,能加速深度模型部署。

调研过程中发现在这块板子上部署模型大多数都是从Keras、TensorFlow开始训练并最终部署,而研究者常用的Pytorch竟然没有教程,于是今天就尝试来讲一讲。

01 使用Pytorch训练分类网络模型

必要软件包安装

pip install tensorbay pillow torch torchvision numpy

数据集获取

一个AccessKey获取所有数据集。

我们使用一个开源数据集平台:https://gas.graviti.com ,这个网站汇总了AI开发者常见的公开数据集,调用其SDK就能直接在线训练,而且许多数据集直接在国内网络下连接直接使用,还是非常方便的。

a. 打开本文对应数据集链接 https://gas.graviti.com/dataset/data-decorators/MNIST

b. 右上角注册登录

c. fork数据集

d. 点击网页上方开发者工具,获取使用SDK所需的AccessKey,获取到 AccessKey 后,将其存在项目根目录的gas_key.py 里:KEY = "<Your-Key>"

通过AccessKey可以上传数据、读取数据、使用数据,灵活对接模型开发和训练,与数据pipeline快速集成。