yolov5 – 导出onnx模型教程

1 yolov5导出onnx模型

yolov5 Github：https://github.com/ultralytics/yolov5

yolov5 导出模型官方文档：https://docs.ultralytics.com/yolov5/tutorials/model_export/

1.1 配置yolov5所需环境

yolov5运行环境需要Python>=3.8，Pytorch>=1.8，使用Anaconda创建一个名为yolov5的虚拟环境，并在该虚拟环境下安装Python>=3.8，Pytorch>=1.8，之后安装yolov5的相关依赖。

拉取项目，安装相关依赖

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
pip install -r requirements.txt  # install

1.2 导出yolov5 onnx模型

先从yolov5的github下载相关模型的预训练模型，这里以yolov5s为例，下载完成之后放在yolov5 git仓库根目录。

切换到yolov5所在的虚拟环境，并在根目录下执行以下命令

python export.py --weights yolov5s.pt --include onnx

命令成功之后，会在根目录下生成yolov5s.onnx的onnx模型文件。

上述导出命令是比较简单的，很有可能在执行上述命令的过程中发生以下错误

ONNX: export failure 0.0s: Unsupported ONNX opset version: 17

出现这个错误，是因为不支持onnx opset 17，所以我们需要在导出命令中指定onnx opset version，所以将上述命令修改为

python export.py --weights yolov5s.pt --include onnx --opset 11

即可导出成功。

在export.py中有很多导出模型时可以用到的参数，比如

def parse_opt(known=False):
    parser = argparse.ArgumentParser()
    parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path')
    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)')
    parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)')
    parser.add_argument('--batch-size', type=int, default=1, help='batch size')
    parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
    parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
    parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
    parser.add_argument('--keras', action='store_true', help='TF: use Keras')
    parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
    parser.add_argument('--int8', action='store_true', help='CoreML/TF/OpenVINO INT8 quantization')
    parser.add_argument('--dynamic', action='store_true', help='ONNX/TF/TensorRT: dynamic axes')
    parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model')
    parser.add_argument('--opset', type=int, default=17, help='ONNX: opset version')
    parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
    parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
    parser.add_argument('--nms', action='store_true', help='TF: add NMS to model')
    parser.add_argument('--agnostic-nms', action='store_true', help='TF: add agnostic NMS to model')
    parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')
    parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep')
    parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold')
    parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold')
    parser.add_argument(
        '--include',
        nargs='+',
        default=['torchscript'],
        help='torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle')
    opt = parser.parse_known_args()[0] if known else parser.parse_args()
    print_args(vars(opt))
    return opt

根据上述的导出参数，我们可以自定义导出onnx模型。

比如我们可以导出fp16半精度的yolov5s的onnx模型，命令行如下

python export.py --weights yolov5s.pt --include onnx --opset 11 --half --device 0

默认导出的yolov5s模型时支持图片640x640，我们可以通过以下命令将输入图片的大小修改为320x320

python export.py --weights yolov5s.pt --include onnx --img-size 320 320

在导出时简化模型

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

更多的命令参数作用可以慢慢去摸索和扩展。

参考

• https://docs.ultralytics.com/yolov5/tutorials/model_export/