yolov5导出onnx失败
yolov5
yolov5 - Ultralytics YOLOv8的前身,是一个用于目标检测、图像分割和图像分类任务的先进模型。
项目地址:https://gitcode.com/gh_mirrors/yo/yolov5
免费下载资源
·
最近看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的最新版本也可以解决问题。
GitHub 加速计划 / yo / yolov5
49.43 K
16.03 K
下载
yolov5 - Ultralytics YOLOv8的前身,是一个用于目标检测、图像分割和图像分类任务的先进模型。
最近提交(Master分支:3 个月前 )
79b7336f
* Update Integrations table
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
* Update README.md
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
* Update README.zh-CN.md
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
---------
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com> 1 个月前
94a62456
* fix: quad training
* fix: quad training in segmentation 1 个月前
更多推荐
已为社区贡献15条内容
所有评论(0)