AX620A运行yolov5s自训练模型全过程记录（windows）

AX620A运行yolov5s自训练模型全过程记录（windows）

---

前言

上一篇文章《ubuntu22.04搭建AX620A官方例程开发环境》记录了AX620A开发环境的搭建，这段时间板子终于到手了，尝试了一下怎么跑自己训练的yolov5模型，将整个过程记录一下。

一、搭建模型转换环境（使用GPU加速）

  模型转换工具名称叫做Pulsar工具链，官方提供了docker容器环境。因为我十分懒得装一台ubuntu实体机，所以一切操作都想办法在windows下实现的~~~

1. 安装Windows平台的Docker-Desktop客户端
   下载地址：点这里
   无脑安装就行了，没有什么特别操作
   安装好后，打开Docker-Desktop，程序会初始化和启动，然后就可以打开cmd输入命令拉取工具链镜像了

docker pull sipeed/pulsar

2. 下载镜像成功后，应该能在Docker-Desktop的images列表中看到了
   
3. 原本到这里就可以直接启动容器来转换模型了，不过突然看到工具链是支持gpu加速的，但镜像里面没有安装相应的cuda环境，所以这里就要一番折腾了。
   思路是：安装WSL2——安装ubuntu子系统——子系统内安装nVidia CUDA toolkit——再安装nvidia-docker。这一番操作后，以后在windows cmd下启动的容器（需加上“--gpus all”参数），就都带上完整的cuda环境了。

• 使用WSL2安装ubuntu20.04
    恕我一直孤陋寡闻，直到这时才知道原来windows下有这么个方便的东西，以前一直只会装虚拟机…
  WSL2具体怎么开启百度一下，反正我的win10更新到最新版本操作没遇到什么问题。

(1) 查看windows平台可支持的linux子系统列表

wsl --list --online
以下是可安装的有效分发的列表。
请使用“wsl --install -d <分发>”安装。

NAME            FRIENDLY NAME
Ubuntu          Ubuntu
Debian          Debian GNU/Linux
kali-linux      Kali Linux Rolling
openSUSE-42     openSUSE Leap 42
SLES-12         SUSE Linux Enterprise Server v12
Ubuntu-16.04    Ubuntu 16.04 LTS
Ubuntu-18.04    Ubuntu 18.04 LTS
Ubuntu-20.04    Ubuntu 20.04 LTS

(2) 设置每个安装的发行版默认用WSL2启动

$ wsl --set-default-version 2

(3) 安装Ubuntu20.04

$ wsl --install -d  Ubuntu-20.04

(4) 设置默认子系统为Ubuntu20.04

$ wsl --setdefault Ubuntu-20.04

(5) 查看子系统信息，星号为缺省系统

$ wsl -l -v
  NAME                   STATE           VERSION
* Ubuntu-20.04           Stopped         2
  docker-desktop-data    Running         2
  docker-desktop         Running         2

• ubuntu22.04下安装cuda
  以下操作均在windows的子系统ubuntu终端输入，逐条执行

wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin

sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600

wget https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda-repo-wsl-ubuntu-11-7-local_11.7.0-1_amd64.deb

sudo dpkg -i cuda-repo-wsl-ubuntu-11-7-local_11.7.0-1_amd64.deb

sudo cp /var/cuda-repo-wsl-ubuntu-11-7-local/cuda-B81839D3-keyring.gpg /usr/share/keyrings/

sudo apt-get update

sudo apt-get -y install cuda

• ubuntu22.04下安装nvidia-docker

distribution=$(. /etc/os-release;echo $ID$VERSION_ID)

curl -s -L https://nvidia.github.io/nvidia-docker/gpgkey | sudo apt-key add -

curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.list | sudo tee 
/etc/apt/sources.list.d/nvidia-docker.list

sudo apt-get update

sudo apt-get install -y nvidia-docker2

至此，cuda的相关环境就安装好了，可以在ubuntu的终端下输入：nvidia-smi确认显卡是否有识别到

• 启动容器
  在windows的cmd终端中，输入以下命令来启动pulsar容器

docker run -it --net host  --gpus all --shm-size 12g -v D:\ax620_data:/data sipeed/pulsar --name ax620

其中：
--gpus all要加上，容器才能使用gpu。
--shm-size 12g 指定容器使用内存，根据你实际内存大小来确定。我这个数值转换yolov5s的模型没有问题。
--name ax620指定容器名称，以后在cmd下输入:docker start ax620 && docker attach ax620就可以方便进入容器环境内。
D:\ax620_data:/data 将本地硬盘指定路径映射到容器的/data目录。这个文件夹需要存放模型转换用到的工程，可以到这里下载

二、转换yolov5自训练模型

1. yolov5使用的是最新的V6.2版，按照标准流程训练出来pt模型文件后，使用export.py导出为onnx格式

python export.py --weights yourModel.pt --simplify --include onnx

2. 要将onnx三个输出层的后处理去掉，因为这部分的处理将在ax620程序里面实现处理。可参考这位大哥的文章《爱芯元智AX620A部署yolov5 6.0模型实录》。具体处理的python代码如下：

import onnx

input_path = r"D:\ax620_data\model\yourModel.onnx"
output_path = r"D:\ax620_data\model\yourModel_rip.onnx"
input_names = [