技术博客丨我用深度学习做个视觉AI微型处理器

作者：张强，Datawhale成员

讲多了算法，如何真正将算法应用到产品领域？本文将带你从0用深度学习打造一个视觉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竟然没有教程，于是今天就尝试来讲一讲。

接下来，我们将从使用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快速集成。

e. AccessKey写入后就可以写代码读取数据了，读取后可以使用一行代码自行下载，或者可以开启缓存功能，在读取过后会自动将数据存储到本地。将下载后的数据放在data文件夹下：

import numpy as np
from PIL import Image

from tensorbay import GAS
from tensorbay.dataset import Dataset
from tensorbay.dataset import Segment

def read_gas_image(data):
    with data.open() as fp:
        image = Image.open(fp)
    return np.array(image)
  
KEY = "用你的Key替换掉这个字符串"
# Authorize a GAS client.
gas = GAS(KEY)
# Get a dataset.
dataset = Dataset("MNIST", gas)

# 开启下行语句在当前路径下的data目录缓存数据
# dataset.enable_cache("data")

# List dataset segments.
segments = dataset.keys()
# Get a segment by name
segment = dataset["train"]
for data in segment:
    # 图片数据
    image = read_gas_image(data)
    # 标签数据
    label = data.label.classification.category

怎么把这个数据集集成到Pytorch里呢？官方文档也为此写了不少例子[4]。笔者尝试过觉得挺方便，在为不同任务训练嵌入式AI模型时，只需更换数据集的名字，就能集成，不用再打开浏览器、等待下载以及处理很久了。有了数据集之后，我们接下来用其训练一个分类任务模型。

深度网络模型选型

结合硬件特点设计网络。

在考虑硬件部署的任务时，网络的设计就要受到些许限制。

首先，大而深的模型是不行的，K210的RAM是6M，这意味着模型+程序都要烧到这个空间。当然我们可以放到内存卡中，但实时性要受影响。

其次，还要考虑AI编译器对特定算子的优化，以K210 NNCASE为例[3]，其支持TFLite、Caffe、ONNX共三个平台的算子。

打开对应平台，能够到具体有哪些算子已经实现了低层优化。可以看到对ONNX算子优化还是比较多的。如果所选用的网络算子较新，抑或是模型太大，都要在本步多加思考与设计。

如果如果最后部署不成功，往往需要回到这一步考虑网络的设计。为了尽可能减少算子的使用，本文设计一个只基于卷积+ReLU+Pool的CNN：

代码文件名：models/net.py

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(1, 6, 5)
        self.relu1 = nn.ReLU()
        self.pool1 = nn.MaxPool2d(2)

        self.conv2 = nn.Conv2d(6, 16, 5)
        self.relu2 = nn.ReLU()
        self.pool2 = nn.MaxPool2d(2)

        self.conv3 = nn.Conv2d(16, 32, 4)
        self.relu3 = nn.ReLU()

        self.conv4 = nn.Conv2d(32, 64, 1)
        self.relu4 = nn.ReLU()

        self.conv5 = nn.Conv2d(64, 32, 1)
        self.relu5 = nn.ReLU()

        self.conv6 = nn.Conv2d(32, 10, 1)
        self.relu6 = nn.ReLU()

    def forward(self, x):
        y = self.conv1(x)
        y = self.relu1(y)
        y = self.pool1(y)
        y = self.conv2(y)
        y = self.relu2(y)
        y = self.pool2(y)
        y = self.conv3(y)
        y = self.relu3(y)
        y = self.conv4(y)
        y = self.relu4(y)
        y = self.conv5(y)
        y = self.relu6(y)
        y = self.conv6(y)
        y = self.relu6(y)

        y = y.view(y.shape[0], -1)

        return y

网络训练

设计好模型后，使用如下脚本进行训练。接下来脚本文件大致浏览一下，明白其中的工作原理即可。

代码文件名：1.train.py

注意将其中的ACCESS_KEY替成你自己的AccessKey。

from __future__ import print_function
import argparse
import torch
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import StepLR
from models.net import Net
from PIL import Image
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

from tensorbay import GAS
from tensorbay.dataset import Dataset as TensorBayDataset


class MNISTSegment(Dataset):
    """class for wrapping a MNIST segment."""

    def __init__(self, gas, segment_name, transform, cache=True):
        super().__init__()
        self.dataset = TensorBayDataset("MNIST", gas)
        if cache:
            self.dataset.enable_cache("data")
        self.segment = self.dataset[segment_name]