PyTorch模型部署到Android

YOLOV5模型部署到Android的过程记录，系统环境：Ubuntu

1. PyTorch模型转ONNX

cd path/to/yolov5
python models/export.py --weights runs/train/exp/weights/best.pt --train --include onnx

其中第二行命令中的--weights指保存的参数文件；

--train表示模型以train模式进行，主要是根据yolov5的代码来看，需要导出train模式的网络结构，因为forward需要return三个输出，如果是eval模式，输出的不是这种格式；

--include表示需要导出的目标格式有哪些，因为目前只需要onnx，所以需要只写onnx，默认导出各种支持的格式

cd path/to/yolov5/runs/train/exp/weights
python -m onnxsim best.onnx best-sim.onnx

其中模块onnxsim是需要下载的，使用pip install onnx-simplifier；安装完onnx-simplifier后即可使用如上命令来简化onnx文件，得到best-sim.onnx；

可以到github下载ncnn源码并自行编译，本文只讲预编译版本的使用。

我的环境是Ubuntu16，所以下载的是ncnn-20210525-ubuntu-1604-shared.zip 。目前不清楚shared版和无shared版的区别，不过貌似不带shared的版本缺一些东西导致报错。

将压缩文件解压后即可使用ncnn了。

cd path/to/yolov5/runs/train/exp/weights
path/to/ncnn/bin/onnx2ncnn best-sim.onnx yolov5s-traffic.param yolov5s-traffic.bin

不出所料，转化完成，但是出现如下报错：

转换过程中有一些不支持的结构

此时需要手动修改生成的param文件来解决错误。

下图为param文件修改前的情况：

修改前的内容

转换错误的层为5到12行的Crop操作，因此需要将4到13行全部删除并体会为自定义层YoloV5Focus。因为一共减少了9层，所以将第2行的第1个值减9：

换为自定义层后的内容

另外，网络最后有三个reshape层，我们需要实现输入任意尺寸的图片，因此需要将位置0的大小改为-1以实现自动计算大小：

修改reshape层的内容

path/to/ncnn/bin/ncnnoptimize yolov5s-traffic.param yolov5s-traffic.bin yolov5s-traffic-opt.param yolov5s-traffic-opt.bin 1

其中末尾的1表示将float32优化为float16以降低文件大小。