yolov5导出onnx失败

        最近看yolov5更新了代码,支持导出的模型增加了不少,clone了最新的代码,但是遇到了onnx模型导出失败的问题

        最新的代码是v6.1的:

        导出模型时提示:

        然后就没了下文,直接退出了,使用旧版的export.py作为包导入export_onnx()也一样报错,看了下代码,代码写的其实是一模一样的,只是打印信息使用了log,但之前的代码导出onnx是没问题的,这就奇怪了。

        折腾了半天,差点直接重装anaconda。。。github上有人是这么解决的。。。

        最后发现是导包的语句在新版本中放的位置不对

import onnx
def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')):
    # YOLOv5 ONNX export
    try:
        check_requirements(('onnx',))
        # 将以下导包语句放到函数外部
        # import onnx

        LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
        f = file.with_suffix('.onnx')

        torch.onnx.export(model, im, f, verbose=False, opset_version=opset,
                          training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
                          do_constant_folding=not train,
                          input_names=['images'],
                          output_names=['output'],
                          dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)
                                        'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
                                        } if dynamic else None)

        # Checks
        model_onnx = onnx.load(f)  # load onnx model
        onnx.checker.check_model(model_onnx)  # check onnx model
        # LOGGER.info(onnx.helper.printable_graph(model_onnx.graph))  # print

        # Simplify
        if simplify:
            try:
                check_requirements(('onnx-simplifier',))
                import onnxsim

                LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
                model_onnx, check = onnxsim.simplify(
                    model_onnx,
                    dynamic_input_shape=dynamic,
                    input_shapes={'images': list(im.shape)} if dynamic else None)
                assert check, 'assert check failed'
                onnx.save(model_onnx, f)
            except Exception as e:
                LOGGER.info(f'{prefix} simplifier failure: {e}')
        LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
        return f
    except Exception as e:
        LOGGER.info(f'{prefix} export failure: {e}')

        将import onnx放到函数的外部,问题解决

        另一个方法就是卸载当前的onnx,pip安装onnx的最新版本也可以解决问题。

猜你喜欢

转载自blog.csdn.net/athrunsunny/article/details/123741320
今日推荐