Commit
•
7926afc
1
Parent(s):
c9a3b14
Add `--half` support for FP16 CoreML exports with (#7446)
Browse files* add fp16 for coreml using --half
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* fix
* Cleanup
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
* Update export.py
* [pre-commit.ci] auto fixes from pre-commit.com hooks
for more information, see https://pre-commit.ci
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
export.py
CHANGED
@@ -186,7 +186,7 @@ def export_openvino(model, im, file, prefix=colorstr('OpenVINO:')):
|
|
186 |
LOGGER.info(f'\n{prefix} export failure: {e}')
|
187 |
|
188 |
|
189 |
-
def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
|
190 |
# YOLOv5 CoreML export
|
191 |
try:
|
192 |
check_requirements(('coremltools',))
|
@@ -197,6 +197,14 @@ def export_coreml(model, im, file, prefix=colorstr('CoreML:')):
|
|
197 |
|
198 |
ts = torch.jit.trace(model, im, strict=False) # TorchScript model
|
199 |
ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
200 |
ct_model.save(f)
|
201 |
|
202 |
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
@@ -466,7 +474,8 @@ def run(
|
|
466 |
|
467 |
# Load PyTorch model
|
468 |
device = select_device(device)
|
469 |
-
|
|
|
470 |
model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
|
471 |
nc, names = model.nc, model.names # number of classes, class names
|
472 |
|
@@ -480,7 +489,7 @@ def run(
|
|
480 |
im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection
|
481 |
|
482 |
# Update model
|
483 |
-
if half:
|
484 |
im, model = im.half(), model.half() # to FP16
|
485 |
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
486 |
for k, m in model.named_modules():
|
@@ -506,7 +515,7 @@ def run(
|
|
506 |
if xml: # OpenVINO
|
507 |
f[3] = export_openvino(model, im, file)
|
508 |
if coreml:
|
509 |
-
_, f[4] = export_coreml(model, im, file)
|
510 |
|
511 |
# TensorFlow Exports
|
512 |
if any((saved_model, pb, tflite, edgetpu, tfjs)):
|
|
|
186 |
LOGGER.info(f'\n{prefix} export failure: {e}')
|
187 |
|
188 |
|
189 |
+
def export_coreml(model, im, file, int8, half, prefix=colorstr('CoreML:')):
|
190 |
# YOLOv5 CoreML export
|
191 |
try:
|
192 |
check_requirements(('coremltools',))
|
|
|
197 |
|
198 |
ts = torch.jit.trace(model, im, strict=False) # TorchScript model
|
199 |
ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])])
|
200 |
+
bits, mode = (8, 'kmeans_lut') if int8 else (16, 'linear') if half else (32, None)
|
201 |
+
if bits < 32:
|
202 |
+
if platform.system() == 'Darwin': # quantization only supported on macOS
|
203 |
+
with warnings.catch_warnings():
|
204 |
+
warnings.filterwarnings("ignore", category=DeprecationWarning) # suppress numpy==1.20 float warning
|
205 |
+
ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode)
|
206 |
+
else:
|
207 |
+
print(f'{prefix} quantization only supported on macOS, skipping...')
|
208 |
ct_model.save(f)
|
209 |
|
210 |
LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)')
|
|
|
474 |
|
475 |
# Load PyTorch model
|
476 |
device = select_device(device)
|
477 |
+
if half:
|
478 |
+
assert device.type != 'cpu' or coreml, '--half only compatible with GPU export, i.e. use --device 0'
|
479 |
model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model
|
480 |
nc, names = model.nc, model.names # number of classes, class names
|
481 |
|
|
|
489 |
im = torch.zeros(batch_size, 3, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection
|
490 |
|
491 |
# Update model
|
492 |
+
if half and not coreml:
|
493 |
im, model = im.half(), model.half() # to FP16
|
494 |
model.train() if train else model.eval() # training mode = no Detect() layer grid construction
|
495 |
for k, m in model.named_modules():
|
|
|
515 |
if xml: # OpenVINO
|
516 |
f[3] = export_openvino(model, im, file)
|
517 |
if coreml:
|
518 |
+
_, f[4] = export_coreml(model, im, file, int8, half)
|
519 |
|
520 |
# TensorFlow Exports
|
521 |
if any((saved_model, pb, tflite, edgetpu, tfjs)):
|