Skip to content

Y
yolov5
Unverified 提交 d9b4e6b7 authored 作者: Cristi Fati's avatar Cristi Fati 提交者： GitHub
浏览文件
操作

Add `--include torchscript onnx coreml` argument (#3137)

* Allow users to skip exporting in formats that they don't care about * Correct comments * Update export.py renamed --skip-format to --exclude * Switched format from exclude to include (as instructed by @glenn-jocher) * cleanup Co-authored-by: 's avatarGlenn Jocher <glenn.jocher@ultralytics.com>
上级 bd6f6a78
隐藏空白字符变更
内嵌 并排
正在显示 包含 61 行增加 和 55 行删除
models/export.py
浏览文件 @ d9b4e6b7
"""Exports a YOLOv5 *.pt model to ONNX and TorchScript formats """Exports a YOLOv5 *.pt model to TorchScript, ONNX, CoreML formats
Usage: Usage:
$ export PYTHONPATH="$PWD" && python models/export.py --weights yolov5s.pt --img 640 --batch 1 $ python path/to/models/export.py --weights yolov5s.pt --img 640 --batch 1
""" """
import argparse import argparse
...@@ -27,6 +27,7 @@ if __name__ == '__main__': ...@@ -27,6 +27,7 @@ if __name__ == '__main__':
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='image size') # height, width
parser.add_argument('--batch-size', type=int, default=1, help='batch size') parser.add_argument('--batch-size', type=int, default=1, help='batch size')
parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--include', nargs='+', default=['torchscript', 'onnx', 'coreml'], help='include formats')
parser.add_argument('--half', action='store_true', help='FP16 half-precision export') parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True') parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
parser.add_argument('--train', action='store_true', help='model.train() mode') parser.add_argument('--train', action='store_true', help='model.train() mode')
...@@ -35,6 +36,7 @@ if __name__ == '__main__': ...@@ -35,6 +36,7 @@ if __name__ == '__main__':
parser.add_argument('--simplify', action='store_true', help='simplify ONNX model') # ONNX-only parser.add_argument('--simplify', action='store_true', help='simplify ONNX model') # ONNX-only
opt = parser.parse_args() opt = parser.parse_args()
opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand opt.img_size *= 2 if len(opt.img_size) == 1 else 1 # expand
opt.include = [x.lower() for x in opt.include]
print(opt) print(opt)
set_logging() set_logging()
t = time.time() t = time.time()
...@@ -47,7 +49,7 @@ if __name__ == '__main__': ...@@ -47,7 +49,7 @@ if __name__ == '__main__':
# Checks # Checks
gs = int(max(model.stride)) # grid size (max stride) gs = int(max(model.stride)) # grid size (max stride)
opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples opt.img_size = [check_img_size(x, gs) for x in opt.img_size] # verify img_size are gs-multiples
assert not (opt.device.lower() == "cpu" and opt.half), '--half only compatible with GPU export, i.e. use --device 0' assert not (opt.device.lower() == 'cpu' and opt.half), '--half only compatible with GPU export, i.e. use --device 0'
# Input # Input
img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device) # image size(1,3,320,192) iDetection img = torch.zeros(opt.batch_size, 3, *opt.img_size).to(device) # image size(1,3,320,192) iDetection
...@@ -74,62 +76,66 @@ if __name__ == '__main__': ...@@ -74,62 +76,66 @@ if __name__ == '__main__':
print(f"\n{colorstr('PyTorch:')} starting from {opt.weights} ({file_size(opt.weights):.1f} MB)") print(f"\n{colorstr('PyTorch:')} starting from {opt.weights} ({file_size(opt.weights):.1f} MB)")
# TorchScript export ----------------------------------------------------------------------------------------------- # TorchScript export -----------------------------------------------------------------------------------------------
prefix = colorstr('TorchScript:') if 'torchscript' in opt.include or 'coreml' in opt.include:
try: prefix = colorstr('TorchScript:')
print(f'\n{prefix} starting export with torch {torch.__version__}...') try:
f = opt.weights.replace('.pt', '.torchscript.pt') # filename print(f'\n{prefix} starting export with torch {torch.__version__}...')
ts = torch.jit.trace(model, img, strict=False) f = opt.weights.replace('.pt', '.torchscript.pt') # filename
(optimize_for_mobile(ts) if opt.optimize else ts).save(f) ts = torch.jit.trace(model, img, strict=False)
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') (optimize_for_mobile(ts) if opt.optimize else ts).save(f)
except Exception as e: print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
print(f'{prefix} export failure: {e}') except Exception as e:
print(f'{prefix} export failure: {e}')
# ONNX export ------------------------------------------------------------------------------------------------------ # ONNX export ------------------------------------------------------------------------------------------------------
prefix = colorstr('ONNX:') if 'onnx' in opt.include:
try: prefix = colorstr('ONNX:')
import onnx try:
import onnx
print(f'{prefix} starting export with onnx {onnx.__version__}...')
f = opt.weights.replace('.pt', '.onnx') # filename print(f'{prefix} starting export with onnx {onnx.__version__}...')
torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'], f = opt.weights.replace('.pt', '.onnx') # filename
dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640) torch.onnx.export(model, img, f, verbose=False, opset_version=12, input_names=['images'],
'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None) dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # size(1,3,640,640)
'output': {0: 'batch', 2: 'y', 3: 'x'}} if opt.dynamic else None)
# Checks
model_onnx = onnx.load(f) # load onnx model # Checks
onnx.checker.check_model(model_onnx) # check onnx model model_onnx = onnx.load(f) # load onnx model
# print(onnx.helper.printable_graph(model_onnx.graph)) # print onnx.checker.check_model(model_onnx) # check onnx model
# print(onnx.helper.printable_graph(model_onnx.graph)) # print
# Simplify
if opt.simplify: # Simplify
try: if opt.simplify:
check_requirements(['onnx-simplifier']) try:
import onnxsim check_requirements(['onnx-simplifier'])
import onnxsim
print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
model_onnx, check = onnxsim.simplify(model_onnx, print(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
dynamic_input_shape=opt.dynamic, model_onnx, check = onnxsim.simplify(
input_shapes={'images': list(img.shape)} if opt.dynamic else None) model_onnx,
assert check, 'assert check failed' dynamic_input_shape=opt.dynamic,
onnx.save(model_onnx, f) input_shapes={'images': list(img.shape)} if opt.dynamic else None)
except Exception as e: assert check, 'assert check failed'
print(f'{prefix} simplifier failure: {e}') onnx.save(model_onnx, f)
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') except Exception as e:
except Exception as e: print(f'{prefix} simplifier failure: {e}')
print(f'{prefix} export failure: {e}') print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
except Exception as e:
print(f'{prefix} export failure: {e}')
# CoreML export ---------------------------------------------------------------------------------------------------- # CoreML export ----------------------------------------------------------------------------------------------------
prefix = colorstr('CoreML:') if 'coreml' in opt.include:
try: prefix = colorstr('CoreML:')
import coremltools as ct try:
import coremltools as ct
print(f'{prefix} starting export with coremltools {ct.__version__}...')
model = ct.convert(ts, inputs=[ct.ImageType(name='image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])]) print(f'{prefix} starting export with coremltools {ct.__version__}...')
f = opt.weights.replace('.pt', '.mlmodel') # filename model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
model.save(f) f = opt.weights.replace('.pt', '.mlmodel') # filename
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') model.save(f)
except Exception as e: print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
print(f'{prefix} export failure: {e}') except Exception as e:
print(f'{prefix} export failure: {e}')
# Finish # Finish
print(f'\nExport complete ({time.time() - t:.2f}s). Visualize with https://github.com/lutzroeder/netron.') print(f'\nExport complete ({time.time() - t:.2f}s). Visualize with https://github.com/lutzroeder/netron.')
Markdown 格式
0% 或
您添加了 0 人 到此讨论。请谨慎行事。
请先完成此评论的编辑！
请 注册 或者 后发表评论