用定制开发板通过vitis ai 2.0部署自己训练的yolov3（pytorch框架）

前言

开发板为alinx黑金xczu9eg，具体配置教程请看我的另一篇文章

pc端：yolov3模型为官方ulteralytics-yolov3-9.5.0

pytorch版本为1.7.1

根据vitis ai官网要求，接下来第一章和第二章介绍如何准备这三个文件

第三第四章介绍部署到fpga的具体细节

一、训练自己的yolov3模型

1.在github的ulteralytics官网下载模型文件GitHub - ultralytics/yolov3: YOLOv3 in PyTorch > ONNX > CoreML > TFLite

2.修改模型训练文件train.py，从他原本只保存模型权重字典，改为用torch.save（）保存模型结构＋权重的.pth文件

​

注意： 此处模型输入图像的尺寸应当与dpu尺寸相同，我的是416*416

​

由此得到model.pth文件 

二、准备模型文件和数据集文件

1.模型文件

模型文件yolo.py在官方的项目工程中

​

并且要对其更改，确保仅包含前传函数

具体参考后面大佬提供的vitis ai工程 的yolov3模型定义文件

2.数据集文件

搞几百张训练时的数据集放在文件夹

三、量化编译生成.xmodel

下载github大佬上传的 vitis ai例程

GitHub - HSqure/ultralytics-pt-yolov3-vitis-ai-edge: This demo is only used for inference testing of Vitis AI v1.4 and quantitative compilation of DPU. It is compatible with the training results of v9.5.0 version of ultramatics (it needs to use the model saving method of Python 1.4 version)

在虚拟机中进入vitis ai pytorch的环境，然后 

​

将之前准备的数据集（图片）放在 test_sample/val文件夹下

浮点模型定义文件yolo.py放在models文件夹下，

训练好的模型文件改名为yolov3-pedestrian.pth，放在weights文件夹下

​

之后再工程文件的目录下运行

python quant.py --quant_mode calib --subset_len 1

​

然后运行 

python quantize.py --quant_mode test --subset_len 1 --batch_size=1 --deploy

​

然后准备自己板子硬件配置时候的arch.json文件

运行以下命令

vai_c_xir -x quantize_result/Model_int.xmodel -a arch.json -o xmodeloutput -n cocoxmodel

​

生成最终的 .xmodel文件

四、准备可执行文件和.prototxt文件

这一步开始准备上板运行所需的文件

1.执行文件

首先在下载下来的vitis ai2.0官方的Vitis-AI-2.0/demo/Vitis-AI-Library/samples/yolov3 文件夹下，运行以下命令通过个人配置的路径，进入以下环境

source  <path to petalinux_sdk>/environment-setup-cortexa72-cortexa53-xilinx-linux

然后通过

bash -x build.sh

生成可执行文件

​

2. .prototxt文件

修改官方示例中的.prototxt文件

model {
  kernel {
     mean: 0.0
     mean: 0.0
     mean: 0.0
     scale: 0.00390625
     scale: 0.00390625
     scale: 0.00390625
  }
  model_type : YOLOv3
  yolo_v3_param {
    num_classes: 80
    anchorCnt: 3
    layer_name: "Model__Model_Detect_model__Detect_28__Conv2d_m__ModuleList_1__12942_fix"
    layer_name: "Model__Model_Detect_model__Detect_28__Conv2d_m__ModuleList_2__12989_fix"
    layer_name: "Model__Model_Detect_model__Detect_28__Conv2d_m__ModuleList_0__12895_fix"
    conf_threshold: 0.0001
    nms_threshold: 0.005
    biases: 10
    biases: 13
    biases: 16
    biases: 30
    biases: 33
    biases: 23
    biases: 30
    biases: 61
    biases: 62
    biases: 45
    biases: 59
    biases: 119
    biases: 116
    biases: 90
    biases: 156
    biases: 198
    biases: 373
    biases: 326
    test_mAP: false
  }
}

其中，yolov3的这三个输出头名称一定要改成自己的输出层名字

​ 输出层的名字如何确定呢

打开开发板后，通过

xdputil xmodel coco/cocoxmodel.xmodel -l

 打印之前的xmodel信息

​

在最下面有模型名称，请在prototxt文件中更换为下面的模型名称，否则会造成segmentation fault 

​

接下来，将模型的.xmodel和.prototxt文件，还有之前的可执行文件传到板子，通过之前常用的命令就可以执行啦。

​

目前效果还有缺陷， 后续会继续优化

---

总结

本文介绍如何用定制开发板通过vitis ai 2.0部署自己训练的yolov3（pytorch框架）