3. Yolov8 ONNX模型部署源码详解

Yolov8的ONNX转换比较简单，仔细看下源码完全可以明白。

def _export_onnx(self, prefix=colorstr('ONNX:')): # colorstr():改变输出字符串的颜色
    # YOLOv8 ONNX export
    # 对当前环境onnx的版本进行检查
    check_requirements('onnx>=1.12.0')
    import onnx  # noqa

    LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
    # with_suffix(suffix) 有扩展名则替换，无则补充扩展名
    # 转换ONNX文件的存储路径
    f = str(self.file.with_suffix('.onnx'))
    # 模型输出的名字
    output_names = ['output0', 'output1'] if isinstance(self.model, SegmentationModel) else ['output0']
    dynamic = self.args.dynamic
    # 判断是否动态输入和输出
    if dynamic:
        # dynamic 中 'images' 是动态输入的名字，{0: 'batch', 2: 'height', 3: 'width'}是每个动态维度及其名字
        dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}}  # shape(1,3,640,640)
        if isinstance(self.model, SegmentationModel):
            dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
            dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
        elif isinstance(self.model, DetectionModel):
            dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)

    torch.onnx.export(
        # 动态只与CPU兼容
        self.model.cpu() if dynamic else self.model, # 模型 
        self.im.cpu() if dynamic else self.im, # 模型的输入
        f, # 导出ONNX的名字
        verbose=False,
        opset_version=self.args.opset,
        # 使用torch>=1.12的DNN推断可能需要do_constant_folding=False
        do_constant_folding=True,
        input_names=['images'], # 输出张量
        output_names=output_names, # 输入张量
        dynamic_axes=dynamic or None # 设置动态维度
        )

    # Checks
    model_onnx = onnx.load(f)  # load onnx model
    onnx.checker.check_model(model_onnx)  # check onnx model

    # Metadata
    # 自定义元数据填充onnx模型, 存放一些模型相关信息
    d = {'stride': int(max(self.model.stride)), 'names': self.model.names}
    for k, v in d.items():
        meta = model_onnx.metadata_props.add()
        meta.key, meta.value = k, str(v)
    onnx.save(model_onnx, f)