[PETR] 如何将Hailo版本PETR的onnx模型重训练，可视化，导出onnx以及量化编译

PETR是一个统一的框架，用于从多视角图像中进行3D感知的BEV模型。
其模型基本框架如下图：

本文将介绍如何使用Hailo来导出PETR的onnx模型并进行模型编译。
因为Hailo的芯片专注于CNN以及transformer模块的推理，所以此框架中，我们需要将整体模型拆分为Hailo处理部分，以及Host处理部分，请查看下图：

简单来说，将图片处理为2D feature的CNN部分需要hailo上面处理； 然后将这些feature以及其他的3D coordinates 前处理结果一起作为后续 tansformer模型的输入，此transformer部分也在hailo上面运行；但是transformer的Det/Seg Head之类的后处理全部在host端完成。

1. 下载代码

PETR的官方github为：https://github.com/megvii-research/PETR，我们关注的是Hailo版本的github为：https://github.com/hailo-ai/PETR

2. 构建Docker镜像

使用Dockerfile文件创建对应的docker image，参考命令：

docker build --build-arg timezone='cat /etc/timezone' --build-arg user=$USER --build-arg uid=$(id -u) --build-arg gid=$(id -g) --build-arg group=alex -t petr:alex .

3. 运行container

参考命令：

docker run --gpus all --name <docker-name> -it --user $USER --shm-size 32gb -p 5000:5000 '[ -d $HOME/.ssh ] && echo -v $HOME/.ssh:/home/$USER/.ssh' '[ -f $HOME/.gitconfig ] && echo -v $HOME/.gitconfig:/home/$USER/.gitconfig' '[ -f $HOME/.bash_aliases ] && echo -v $HOME/.bash_aliases:/home/$USER/.bash_aliases' -v /home/alex:/home/alex petr:alex /bin/bash --login

4. 激活虚拟环境并准备数据集

source my_venv/bin/activate

数据集准备：

1. 下载并放置nuscenes数据库，可以使用Mini数据集。
   
2. 将文件置于路径：/workspace/PETR/data/nuscenes
3. 在MMDetection3D中实现数据库的生成，参考链接：https://mmdetection3d.readthedocs.io/zh-cn/dev-1.x/advanced_guides/datasets/nuscenes.html
   具体命令：
   cd /workspace/mmdetection3d
   python3 setup.py develop (先安装工具环境)
   python tools/create_data.py nuscenes –root-path /workspace/PETR/data/nuscenes –out-dir /workspace/PETR/data/nuscenes –extra-tag nuscenes
4. 除了标准的mmdetection的数据准备，我们还需要专门为PETR做一些数据准备。
   cd /workspace/PETR/
   python tools/generate_sweep_pkl.py
   在运行前，首先需要修改下这个文档的info_prefix 与 data_root = “data/nuscenes/”
   记得运行两次，分别将info_prefix 修改为“train” 和 “val”

5. 重训练或测试模型

重训练的时候，我们最好是开启虚拟环境 source /workspace/my_venv/bin/activate
并且在这个环境中额外安装 pip install absl-py
# Train
tools/dist_train.sh <config-file-path> <num-gpus> --work-dir <output-workdir>
例如： tools/dist_train.sh projects/configs/petrv2/petrv2_fcos3d_repvgg_b0x32_decoder_3_UN_1600x640.py 1 --work-dir work_dirs
# Test 查看mAP
tools/dist_test.sh <config-file-path> <ckpt-path> <num-gpus> --eval bbox
# 可视化
基本上参考github中的方式
例如：./tools/dist_test.sh projects/configs/petrv2/petrv2_fcos3d_repvgg_b0x32_decoder_3_UN_1600x640.py ckpts/petrv2_fcos3d_repvgg_b0_epoch_120.pth 1 --out work_dirs/pp-nus/results_eval.pkl --format-only --eval-options 'jsonfile_prefix=work_dirs/pp-nus/results_eval'
这里特定的pth对应的是特定的模型，请确保config-file-path中的 num_query 为900.

最后调用create_visualization_video.py将可视化的结果export到一个mp4

python3 tools/visualize.py
python3 create_visualization_video.py
注：需要提前修改 visualize.py中的database路径.

6. 导出模型

将统一框架的模型进行拆分

python tools/export_onnx.py projects/configs/petrv2/petrv2_fcos3d_repvgg_b0x32_decoder_3_UN_1600x640.py petrv2_fcos3d_repvgg_b0_epoch_120.pth --split transformer --ncams 4 --out petrv2_transformer.onnx --img-dims 608 800
python tools/export_onnx.py projects/configs/petrv2/petrv2_fcos3d_repvgg_b0x32_decoder_3_UN_1600x640.py petrv2_fcos3d_repvgg_b0_epoch_120.pth --split backbone --ncams 4 --out petrv2_backbone.onnx --img-dims 608 800

需要将transformer的模型进行剪枝整理，方便hailo对它进行支持。

python tools/onnx_reshape.py --onnx petrv2_transformer.onnx --cfg tools/onnx_reshape_cfg_repvgg_b0x32_cams_4_decoder_3_UN_800x608.json --out petrv2_transformer_const0.onnx

7. 模型量化编译

按照正常的DFC模型进行量化，主要要使用transfomer的专用alls文件进行量化以及compiler。

以上就是使用Hailo利用PETR的onnx模型导出以及模型编译的步骤。