源碼解析二 模型轉換 export.py

基於最新yolov5-v6.0

1.重點

2.相關函數

1.重點

一般使用規則：

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

三種格式想要用哪種就要下載相應的包：

1. torchscript 不需要下載對應的包 有Torch就可以
2. onnx: pip install onnx
3. 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'],
    """

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
        ):