Vitis-AI量化编译MNIST手写数字识别-点头深度学习网站

项目背景

在当前的数字化时代，快速准确的手写数字识别技术在各个领域都发挥着越来越重要的作用，尤其是在银行、邮政服务和数字化存档等行业。这些应用常常要在资源受限的环境中运行，例如嵌入式系统或移动设备，这就要求模型不仅要精确，还要高效。Vitis AI 是 Xilinx 提供的一个先进的人工智能推理开发平台，它专门用于在 Xilinx FPGA 和自适应计算加速平台 (ACAP) 设备上部署深度学习模型。利用 Vitis AI 进行模型量化和编译可以充分发掘 AI 加速的全部潜能，同时还保持了良好的精度。这对于需要高速且高效处理大量数据的手写数字识别应用尤为重要。

项目目标

本项目的主要目标是开发一个基于深度学习的手写数字识别系统，并使用 Vitis AI 的量化和编译工具来优化模型，使其适用于性能受限的硬件平台。具体目标包括：

数字识别准确度：提高手写数字识别的准确度，能够准确区分不同风格和书写质量的数字。
模型优化和效率：通过 Vitis AI 对深度学习模型进行量化，将浮点运算转换为低位数的定点运算，以减少计算资源消耗和提高运行速度，使模型适用于资源受限的嵌入式系统和移动设备。
易于集成和部署：使量化和编译后的模型易于集成到各种应用中，并能够在各种不同的硬件平台上顺利部署。

项目应用

该项目有着广泛的应用潜力，可以在以下场景中得到应用：

自动化表单处理：自动识别填写在表单上的手写数字，从而加速数据输入过程，减少手动输入错误。
邮政编码识别：邮政系统中自动识别信件和包裹上的手写邮政编码，提高邮件分类和分发的速度和准确性。
银行支票处理：在银行行业中，用于自动识别支票上的手写数字，加快交易处理速度，同时降低人工审核的需求。
数字化存档：将历史文档中的手写数字自动转换为数字格式，便于存储、检索和分析，特别是在图书馆、档案馆和政府机构中。
边缘计算设备：在资源有限的边缘计算设备上，如智能相机或移动设备中，实现快速且高效的手写数字识别，支持各种实时应用，例如交通监控中的车牌号识别。

数据集描述

本项目使用MNIST 数据集，MNIST 数据库是一个大型手写数字数据库(包含0~9十个数字)，包含 60000 张训练图像和 10000 张测试图像，通常用于训练各种图像处理系统。训练数据集取自美国人口普査局员工，而测试数据集取自美国高中生。所有的手写数字图片的分辨率为28*28。图像以黑白形式表示，数字以黑色呈现在白色背景上。这些图像已经过标准化处理，确保数字位于图像中心，并且大小相同。

数据集官网：http://MNIST handwritten digit database, Yann LeCun, Corinna Cortes and Chris Burges

项目方法

1.环境搭建

Vitis AI 开发环境可在赛灵思硬件平台上加速 AI 推断，包括边缘器件和 Alveo加速器卡。此环境由经过最优化的 IP 核、工具、库、模型和设计示例组成。要使用Vitis AI进行模型量化和编译，首先要进行环境的搭建。环境设置为Ubuntu22.04+Pytorch1.10.0+Vitis ai2.5。在本项目中需要使用vitis ai中的vai_q_pytorch库，有两种安装方法：分别为使用Docker 容器安装和从源代码安装（源代码安装方法见手册https://docs.xilinx.com/r/2.5-%E7%AE%80%E4%BD%93%E4%B8%AD%E6%96%87/ug1414-vitis-ai/PyTorch-%E7%89%88%E6%9C%AC-vai_q_pytorch），在此推荐使用docker容器安装。

Docker 容器是一种轻量级、可移植的、自给自足的软件封装技术，允许开发者将应用程序及其全部依赖项打包在一起，形成一个标准化的单元，用于软件开发、分发和部署。Docker 容器的核心特点包括：

轻量级：Docker 容器共享宿主机的内核，而不需要像虚拟机一样运行一个完整的操作系统。这使得容器更加高效、占用资源更少。
可移植性：由于容器包含了运行应用所需的一切（代码、运行时环境、库、环境变量等），因此可以在任何支持 Docker 的系统上无差别运行。
一致性和隔离：容器为应用提供了隔离的环境，确保应用在开发、测试和生产环境中的一致性和可重复性。
易于管理：Docker 提供了工具和命令行界面，使得容器的创建、部署和管理变得简单。
可扩展性和微服务架构：容器非常适合微服务架构，因为你可以将应用程序分解为一系列较小的、独立的服务，每个服务运行在自己的容器中。

安装流程：

安装虚拟机，下载Ubuntu22.04，在github上下载vitis ai2.5文件包https://github.com/Xilinx/Vitis-AI/tree/v2.5，并将目录切换到此文件夹下。
卸载旧版本docker：

sudo apt-get remove docker docker-engine docker-ce docker.io

添加官方密钥：

sudo apt-get install curl   
curl -fsSL https://download.docker.com/linux/ubuntu/gpg |sudo apt-key add -

检查是否docker是否正常，正常的话可以运行一个hello world进行检查：

systemctl status docker
sudo docker run hello-world

图片[1]-Vitis-AI量化编译MNIST手写数字识别-点头深度学习网站 — 图1 hello-world运行结果

切换用户，上述命令均在root用户下执行，通过切换用户将当前登录的用户添加到 docker 组中，为了让该用户能够执行 Docker 命令而不需要使用sudo，在完成用户切换后需要重启虚拟机。

sudo usermod -aG docker $USER
docker run hello-world #判断切换是否成功

安装vitis ai docker，从官方仓库https://hub.docker.com/r/xilinx/vitis-ai/tags中拉取镜像，先下载后启动。

docker pull xilinx/vitis-ai-cpu:2.5.0.1260
./docker_run.sh xilinx/vitis-ai-cpu:2.5.0.1260
conda activate vitis-ai-pytorch#激活pytorch环境

图片[2]-Vitis-AI量化编译MNIST手写数字识别-点头深度学习网站 — 图2 docker的安装与启动

如果在dacker pull镜像过程中出现下载慢的情况，可以尝试换源：

sudo mkdir -p /etc/docker
sudo tee /etc/docker/daemon.json <<-'EOF'
{
  "registry-mirrors": [
"https://mirror.baidubce.com"
]
}
EOF
sudo systemctl daemon-reload
sudo systemctl restart docker

2.模型训练

本项目通过搭建一个简单的卷积神经网络来对MNIST数据集进行识别，并通过训练得到预训练模型，具体代码如下：

import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F

class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()

        self.network = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=5, stride=2, padding=1),
            nn.BatchNorm2d(16),
            nn.ReLU(inplace=True),
            nn.Conv2d(16, 32, kernel_size=5, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.Conv2d(64, 10, kernel_size=3, stride=3),
            nn.BatchNorm2d(10),
            nn.Flatten()
            )
    def forward(self, x):
        x = self.network(x)
        return x

def train(model, device, train_loader, optimizer, epoch):
    '''
    train the model
    '''
    model.train()
    counter = 0
    print("Epoch "+str(epoch))
    for batch_idx, (data, target) in enumerate(train_loader):
        data, target = data.to(device), target.to(device)
        optimizer.zero_grad()
        x = model(data)
        output = F.log_softmax(input=x,dim=0)
        loss = F.nll_loss(output, target)
        loss.backward()
        optimizer.step()
        counter += 1

def test(model, device, test_loader):
    '''
    test the model
    '''
    model.eval()
    test_loss = 0
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    acc = 100. * correct / len(test_loader.dataset)
    print('\nTest set: Accuracy: {}/{} ({:.2f}%)\n'.format(correct, len(test_loader.dataset), acc))

    return

import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

import argparse
import sys
import os
import shutil

from common import *

torchvision.datasets.MNIST.resources = [
    ('https://ossci-datasets.s3.amazonaws.com/mnist/train-images-idx3-ubyte.gz', 'f68b3c2dcbeaaa9fbdd348bbdeb94873'),
    ('https://ossci-datasets.s3.amazonaws.com/mnist/train-labels-idx1-ubyte.gz', 'd53e105ee54ea40749a09fcbcd1e9432'),
    ('https://ossci-datasets.s3.amazonaws.com/mnist/t10k-images-idx3-ubyte.gz', '9fb629c4189551a2d022fa330f9573f3'),
    ('https://ossci-datasets.s3.amazonaws.com/mnist/t10k-labels-idx1-ubyte.gz', 'ec29112dd5afa0611ce80d1b7f02629c')
]


def train_test(build_dir, batchsize, learnrate, epochs):

    dset_dir = build_dir + '/dataset'
    float_model = build_dir + '/float_model'

    device = torch.device('cpu')

    model = CNN().to(device)
    optimizer = optim.Adam(model.parameters(), lr=learnrate)

    #image datasets
    train_dataset = torchvision.datasets.MNIST(dset_dir, 
                                               train=True, 
                                               download=False,
                                               transform=train_transform)
    test_dataset = torchvision.datasets.MNIST(dset_dir,
                                              train=False, 
                                              download=False,
                                              transform=test_transform)
    #data loaders
    train_loader = torch.utils.data.DataLoader(train_dataset,
                                               batch_size=batchsize, 
                                               shuffle=True)
    test_loader = torch.utils.data.DataLoader(test_dataset,
                                              batch_size=batchsize, 
                                              shuffle=False)

    # training with test after each epoch
    for epoch in range(1, epochs + 1):
        train(model, device, train_loader, optimizer, epoch)
        test(model, device, test_loader)

    # save the trained model
 
    os.makedirs(float_model)   
    save_path = os.path.join(float_model, 'f_model.pth')
    torch.save(model.state_dict(), save_path) 
    print('Trained model written to',save_path)
    return

3.模型量化

模型推理是一个计算密集型流程，需要大量存储器带宽以满足边缘 (Edge) 应用的低时延和高吞吐量要求。通常，进行神经网络训练时，使用的是 32 位浮点权重和激活值。通过将 32 位浮点权重和激活转换为 8 位整数 (INT8) 格式，Vitis AI 量化器即可降低计算复杂性，而不会损失预测精度。定点网络模型所需存储器带宽较少，因此相比浮点模型，速度更快且能效更高。Vitis AI 量化器支持神经网络中存在公用层，包括但不限于卷积、池化、完全连接和批量归一。Vitis AI 量化器当前支持 TensorFlow（1.x 和 2.x）以及 PyTorch。量化器名称分别为 vai_q_tensorflow 和 vai_q_pytorch。

vitis ai总体模型量化流程如图3所示。使用浮点模型作为输入并执行预处理（折叠 batchnorm 并移除推断不需要的节点），然后将权重/偏差和激活量化为给定位宽。整体量化步骤分为两步：第一步是量化校准（calib），第二步进行求值并生成.Xmodel文件。为采集激活统计数据并提升量化模型精度，Vitis AI 量化器必须运行多次推断迭代，以对激活进行校准。因此，校准图像数据集输入是必需的。通常，量化器能够妥善处理100 – 1000张校准图像。量化的代码里面，调用了torch_quantizer来进行量化，量化后的模型一定要用数据集运行一遍（evaluate）。

图片[3]-Vitis-AI量化编译MNIST手写数字识别-点头深度学习网站 — 图3 量化流程

类 torch_quantizer 用于创建量化器对象，具体参数如下所示。

class torch_quantizer(): 
  def __init__(self,
               quant_mode: str, # ['calib', 'test']#量化模式设置
               module: torch.nn.Module,#预训练模型
               input_args: Union[torch.Tensor, Sequence[Any]] = None,#输入张量，形状与要量化的浮点模块的真实输入相同，值可为随机数值
               state_dict_file: Optional[str] = None,
               output_dir: str = "quantize_result",#量化结果和中间文件的目录
               bitwidth: int = 8,#全局量化位宽。默认值为 8
               device: torch.device = torch.device("cuda"),
               quant_config_file: Optional[str] = None):
               ):

量化部分具体代码如下：

import os
import sys
import argparse
import random
import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from pytorch_nndct.apis import torch_quantizer, dump_xmodel
from common import *#模型部分
DIVIDER = '-----------------------------------------'

def quantize(build_dir,quant_mode,batchsize):

  dset_dir = build_dir + '/dataset'#数据集
  float_model = build_dir + '/float_model'#权重文件
  quant_model = build_dir + '/quant_model'#生成的文件

  device = torch.device('cpu')

  # load trained model
  model = CNN().to(device)
  #model.load_state_dict(torch.load(os.path.join(float_model,'f_model.pth')))
  model.load_state_dict(torch.load(os.path.join(float_model, 'f_model.pth'), map_location=device))#模型读取


  optimize = 1

  # override batchsize if in test mode
  if (quant_mode=='test'):
    batchsize = 1 #求值时b为1
  
  rand_in = torch.randn([batchsize, 1, 28, 28])#和mnist数据集具有相同形状的随机数据
  quantizer = torch_quantizer(quant_mode, model, (rand_in), output_dir=quant_model) 
  quantized_model = quantizer.quant_model#获取量化后的模型

  # data loader，使用测试集
  test_dataset = torchvision.datasets.MNIST(dset_dir,
                                            train=False, 
                                            download=True,
                                            transform=test_transform)

  test_loader = torch.utils.data.DataLoader(test_dataset,
                                            batch_size=batchsize, 
                                            shuffle=False)

  # evaluate 
  test(quantized_model, device, test_loader)#使用转换后的模型前传神经网络，在两步量化里面都有

  # export config
  if quant_mode == 'calib':
    quantizer.export_quant_config()#判定张量的量化步骤
  if quant_mode == 'test':
    quantizer.export_xmodel(deploy_check=False, output_dir=quant_model)#对量化模型求值导出xmodel
  
  return

def run_main():

  # construct the argument parser and parse the arguments
  ap = argparse.ArgumentParser()
  ap.add_argument('-d',  '--build_dir',  type=str, default='build',    help='Path to build folder. Default is build')
  ap.add_argument('-q',  '--quant_mode', type=str, default='calib',    choices=['calib','test'], help='Quantization mode (calib or test). Default is calib')#运行要传这个参数
  ap.add_argument('-b',  '--batchsize',  type=int, default=100,        help='Testing batchsize - must be an integer. Default is 100')
  args = ap.parse_args()

  print('\n'+DIVIDER)
  print('PyTorch version : ',torch.__version__)
  print(sys.version)
  print(DIVIDER)
  print(' Command line options:')
  print ('--build_dir    : ',args.build_dir)
  print ('--quant_mode   : ',args.quant_mode)
  print ('--batchsize    : ',args.batchsize)
  print(DIVIDER)

  quantize(args.build_dir,args.quant_mode,args.batchsize)

  return

if __name__ == '__main__':
    run_main()

第一步量化运行命令为：

python quantize.py -q calib

执行此命令会生成一个py文件和一个json文件（如ResNet.py：转换后的 vai_q_pytorch 格式模型；Quant_info.json张量的量化步骤。保留此文件以便对量化模型进行求值。）

第二步量化运行命令为，执行此命令得到.xmodel文件：

python quantize.py -q test -b 1

4.模型编译

Vitis AI 编译器 (VAI_C) 旨在用作为编译器系列的统一接口，此编译器系列专用于对 DPU 系列的神经网络计算执行最优化。每个编译器都能将单个网络模型映射到单一高度优化的DPU指令序列。xilinx中间表示形式 (XIR) 是基于图的 AI 算法中间表示形式，专为在 FPGA 平台上执行 DPU 编译和有效部署而设计。可以直接根据xilinx提供的compiler把xmodel编译成DPU支持的，基于XIR的xmodel。其中的arch.json要和硬件配置文件适配。命令如下：

vai_c_xir -x ./build/quant_model/CNN_int.xmodel -a /opt/vitis_ai/compiler/arch/DPUCZDX8G/KV260/arch.json -o ./ -n my_model

此时即完成了模型的编译，需要注意观察最终出来的DPU subgraph number是不是1，不是1的话请检查模型是不是有DPUhttps://docs.xilinx.com/r/zh-CN/pg338-dpu不支持的OP，在遇到不支持的OP的时候，DPU会分为多个子图执行，由PS处理完后再发送给DPU，拖慢效率。