pytorch版yolov3模型转换onnx的踩坑过程

前言:

首先要明确,yolov3是分为darknet版本和pytorch版本的,pytorch版本早期代码是兼容darknet的网络,但是经过这么长时间的更新,在现在的代码里分为这么两个分支:master和archive

master分支是目前主要更新的版本,伴随着yolov5的大版本更新而不断的更新代码,并且在代码框架内已经不再对darknet进行支持。

archive分支则是对老版本代码的存档,保留了对darknet的支持,但已经不再更新。

具体的内容可以在yolov3-v9.0的readme中查看到,如下图所示。

截至到yolo

1. master分支yolov3导出为onnx模型

由于master版本的yolov3的代码更新已经与yolov5类似,所以作者在文件内置了export.py文件,在该文件下修改对应参数即可完成模型导出。

# 主要修改--data --weights --include参数,如果对opset等参数有需求,也可以进行修改
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 / 'yolov3-tiny.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 YOLOv3 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

或者直接在命令行窗口输入如下命令进行导出,具体细节可参考export.py文件开头的示例进行操作。

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

2. archive分支yolov3导出为onnx模型

相较于最新版的yolov3代码,archive分支的代码更加接近darknet版本的代码,因此导出为onnx模型很麻烦,没有官方的代码来实现,在网上找了很多参考教程都无法实现,基本上能搜到的都试验过,不行。只有一种方法是可行的,那就是pytorch版本的yolov3.pt文件转化为darknet版本的yolov3.weights文件,再从yolov3.weights转化为yolov3.onnx模型。

在archive分支的readme文件内有这么几行命令,意思是内置的代码可以将pt文件和weights互相转换。

利用下方的python代码即可实现pt文件转化为darknet框架的weights文件。

from models import *

# 输入相应的文件地址
convert('cfg/yolov3.cfg', 'weights/yolov3.pt')

转换完成后,根据下方链接代码来进行weights文件转换onnx文件,注意看readme中的指令来修改对应的参数,yolov3、yolov3-spp、yolov3-tiny的参数是不一样的,注意修改。

https://github.com/linghu8812/tensorrt_inference/tree/master/project/Yolov4

至此,完成了老版本的yolov3.pt转换为onnx的全过程

  • 0
    点赞
  • 3
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值