当前位置：网站首页>源码解析二模型转换 export.py

源码解析二模型转换 export.py

2022-08-05 12:48:00 【51CTO】

基于最新yolov5-v6.0

1.重点

一般使用规则：

       
       python export.py --data " " --weights " " --imgsz 512 --simplify --include "onnx"
      
1.

三种格式想要用哪种就要下载相应的包：

torchscript 不需要下载对应的包有Torch就可以
onnx: pip install onnx
coreml: pip install coremltools

2.相关函数

parse_opt():

       
       def parse_opt():
       
    """
       
    data: 数据集目录 默认=ROOT / 'data/coco128.yaml'
       
    weights：权重文件目录 默认=ROOT / 'yolov5s.pt'
       
    img-size: 输入模型的图片size=(height, width) 默认=[640, 640]
       
    batch-size: batch大小 默认=1
       
    device: 模型运行设备 cuda device, i.e. 0 or 0,1,2,3 or cpu 默认=cpu
       
    include: 要将pt文件转为什么格式 可以为单个原始也可以为list 默认=['torchscript', 'onnx', 'coreml']
       
    half: 是否使用半精度FP16export转换 默认=False
       
    inplace: 是否设置 YOLOv5 Detect() inplace=True  默认=False
       
    train: 是否开启model.train() mode 默认=True  coreml转换必须为True
       
    optimize: TorchScript转化参数 是否进行移动端优化  默认=False
       
    int8： 支持CoreML/TF INT8 量化 不支持ONNX
       
    dynamic: ONNX转换参数  dynamic_axes  ONNX转换是否要进行批处理变量  默认=False
       
    simplify: ONNX转换参数 是否简化onnx模型  默认=False
       
    opset: ONNX转换参数 设置ONNX版本  默认=13
       
    topk-per-class: TF.js每一类别都要保留 默认=100
       
    topk-all: TF.js  Topk为所有class保留
       
    iou-thres: TF.js IoU threshold   default=0.45
       
    conf-thres: TF.js  confidence threshold  default=0.25
       
    include:     需要导出的版本  default=['torchscript', 'onnx'],
       
    """
      
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.

def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'

       
       def run(data=ROOT / 'data/coco128.yaml',  # 'dataset.yaml path'
       
        weights=ROOT / 'yolov5s.pt',  # weights path
       
        imgsz=(640, 640),  # image (height, width)
       
        batch_size=1,  # batch size
       
        device='cpu',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
       
        include=('torchscript', 'onnx', 'coreml'),  # include formats
       
        half=False,  # FP16 half-precision export
       
        inplace=False,  # set YOLOv5 Detect() inplace=True
       
        train=False,  # model.train() mode
       
        optimize=False,  # TorchScript: optimize for mobile
       
        int8=False,  # CoreML/TF INT8 quantization
       
        dynamic=False,  # ONNX/TF: dynamic axes
       
        simplify=False,  # ONNX: simplify model
       
        opset=12,  # ONNX: opset version
       
        topk_per_class=100,  # TF.js NMS: topk per class to keep
       
        topk_all=100,  # TF.js NMS: topk for all classes to keep
       
        iou_thres=0.45,  # TF.js NMS: IoU threshold
       
        conf_thres=0.25  # TF.js NMS: confidence threshold
       
        ):
      
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.