Ankle/yolov5_ZWgra
1
0 Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

Addition refactor export.py (#4089)

Browse Source 
* Addition refactor `export.py`

* Update export.py
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-07-20 18:42:27 +02:00
1 parent 442a7abdf2
commit 3bef77f5cb
1 file changed
+10 -9
Download Patch File Download Diff File Expand all files Collapse all files
export.py
+10 -9
View File
@@ -45,7 +45,7 @@ def export_onnx(model, img, file, opset_version, train, dynamic, simplify):
check_requirements(('onnx', 'onnx-simplifier')) check_requirements(('onnx', 'onnx-simplifier'))
import onnx import onnx
print(f'{prefix} starting export with onnx {onnx.__version__}...') print(f'\n{prefix} starting export with onnx {onnx.__version__}...')
f = file.with_suffix('.onnx') f = file.with_suffix('.onnx')
torch.onnx.export(model, img, f, verbose=False, opset_version=opset_version, torch.onnx.export(model, img, f, verbose=False, opset_version=opset_version,
training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL, training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL,
@@ -80,16 +80,17 @@ def export_onnx(model, img, file, opset_version, train, dynamic, simplify):
print(f'{prefix} export failure: {e}') print(f'{prefix} export failure: {e}')
def export_coreml(ts_model, img, file, train): def export_coreml(model, img, file):
# CoreML model export # CoreML model export
prefix = colorstr('CoreML:') prefix = colorstr('CoreML:')
try: try:
import coremltools as ct import coremltools as ct
print(f'{prefix} starting export with coremltools {ct.__version__}...') print(f'\n{prefix} starting export with coremltools {ct.__version__}...')
f = file.with_suffix('.mlmodel') f = file.with_suffix('.mlmodel')
assert train, 'CoreML exports should be placed in model.train() mode with `python export.py --train`' model.train() # CoreML exports should be placed in model.train() mode
model = ct.convert(ts_model, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])]) ts = torch.jit.trace(model, img, strict=False) # TorchScript model
model = ct.convert(ts, inputs=[ct.ImageType('image', shape=img.shape, scale=1 / 255.0, bias=[0, 0, 0])])
model.save(f) model.save(f)
print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') print(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
except Exception as e: except Exception as e:
@@ -145,12 +146,12 @@ def run(weights='./yolov5s.pt', # weights path
print(f"\n{colorstr('PyTorch:')} starting from {weights} ({file_size(weights):.1f} MB)") print(f"\n{colorstr('PyTorch:')} starting from {weights} ({file_size(weights):.1f} MB)")
# Exports # Exports
if 'torchscript' in include:
export_torchscript(model, img, file, optimize)
if 'onnx' in include: if 'onnx' in include:
export_onnx(model, img, file, opset_version, train, dynamic, simplify) export_onnx(model, img, file, opset_version, train, dynamic, simplify)
if 'torchscript' in include or 'coreml' in include: if 'coreml' in include:
ts = export_torchscript(model, img, file, optimize) export_coreml(model, img, file)
if 'coreml' in include:
export_coreml(ts, img, file, train)
# 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.')