@ -9,7 +9,7 @@ TorchScript | `torchscript` | yolov8n.torchscript
ONNX | ` onnx ` | yolov8n . onnx
OpenVINO | ` openvino ` | yolov8n_openvino_model /
TensorRT | ` engine ` | yolov8n . engine
CoreML | ` coreml ` | yolov8n . mlmodel
CoreML | ` coreml ` | yolov8n . mlpackage
TensorFlow SavedModel | ` saved_model ` | yolov8n_saved_model /
TensorFlow GraphDef | ` pb ` | yolov8n . pb
TensorFlow Lite | ` tflite ` | yolov8n . tflite
@ -35,7 +35,7 @@ Inference:
yolov8n . onnx # ONNX Runtime or OpenCV DNN with dnn=True
yolov8n_openvino_model # OpenVINO
yolov8n . engine # TensorRT
yolov8n . mlmodel # CoreML (macOS-only)
yolov8n . mlpackage # CoreML (macOS-only)
yolov8n_saved_model # TensorFlow SavedModel
yolov8n . pb # TensorFlow GraphDef
yolov8n . tflite # TensorFlow Lite
@ -82,7 +82,7 @@ def export_formats():
[ ' ONNX ' , ' onnx ' , ' .onnx ' , True , True ] ,
[ ' OpenVINO ' , ' openvino ' , ' _openvino_model ' , True , False ] ,
[ ' TensorRT ' , ' engine ' , ' .engine ' , False , True ] ,
[ ' CoreML ' , ' coreml ' , ' .mlmodel ' , True , False ] ,
[ ' CoreML ' , ' coreml ' , ' .mlpackage ' , True , False ] ,
[ ' TensorFlow SavedModel ' , ' saved_model ' , ' _saved_model ' , True , True ] ,
[ ' TensorFlow GraphDef ' , ' pb ' , ' .pb ' , True , True ] ,
[ ' TensorFlow Lite ' , ' tflite ' , ' .tflite ' , True , False ] ,
@ -149,8 +149,10 @@ class Exporter:
self . run_callbacks ( ' on_export_start ' )
t = time . time ( )
format = self . args . format . lower ( ) # to lowercase
if format in ( ' tensorrt ' , ' trt ' ) : # engine aliases
if format in ( ' tensorrt ' , ' trt ' ) : # ' engine' aliases
format = ' engine '
if format in ( ' mlmodel ' , ' mlpackage ' , ' mlprogram ' , ' apple ' , ' ios ' ) : # 'coreml' aliases
format = ' coreml '
fmts = tuple ( export_formats ( ) [ ' Argument ' ] [ 1 : ] ) # available export formats
flags = [ x == format for x in fmts ]
if sum ( flags ) != 1 :
@ -319,7 +321,7 @@ class Exporter:
dynamic [ ' output0 ' ] = { 0 : ' batch ' , 2 : ' anchors ' } # shape(1, 84, 8400)
torch . onnx . export (
self . model . cpu ( ) if dynamic else self . model , # -- dynamic only compatible with cpu
self . model . cpu ( ) if dynamic else self . model , # dynamic=True only compatible with cpu
self . im . cpu ( ) if dynamic else self . im ,
f ,
verbose = False ,
@ -461,14 +463,16 @@ class Exporter:
yaml_save ( f / ' metadata.yaml ' , self . metadata ) # add metadata.yaml
return str ( f ) , None
@try_export
def export_coreml ( self , prefix = colorstr ( ' CoreML: ' ) ) :
""" YOLOv8 CoreML export. """
check_requirements ( ' coremltools>=6.0,<=6.2 ' )
mlmodel = self . args . format . lower ( ) == ' mlmodel ' # legacy *.mlmodel export format requested
check_requirements ( ' coremltools>=6.0,<=6.2 ' if mlmodel else ' coremltools>=7.0.b1 ' )
import coremltools as ct # noqa
LOGGER . info ( f ' \n { prefix } starting export with coremltools { ct . __version__ } ... ' )
f = self . file . with_suffix ( ' .mlmodel ' )
f = self . file . with_suffix ( ' .mlmodel ' if mlmodel else ' .mlpackage ' )
if f . is_dir ( ) :
shutil . rmtree ( f )
bias = [ 0.0 , 0.0 , 0.0 ]
scale = 1 / 255
@ -479,20 +483,38 @@ class Exporter:
elif self . model . task == ' detect ' :
model = iOSDetectModel ( self . model , self . im ) if self . args . nms else self . model
else :
# TODO CoreML Segment and Pose model pipelining
if self . args . nms :
LOGGER . warning ( f " { prefix } WARNING ⚠️ ' nms=True ' is only available for Detect models like ' yolov8n.pt ' . " )
# TODO CoreML Segment and Pose model pipelining
model = self . model
ts = torch . jit . trace ( model . eval ( ) , self . im , strict = False ) # TorchScript model
ct_model = ct . convert ( ts ,
inputs = [ ct . ImageType ( ' image ' , shape = self . im . shape , scale = scale , bias = bias ) ] ,
classifier_config = classifier_config )
bits , mode = ( 8 , ' kmeans_lut ' ) if self . args . int8 else ( 16 , ' linear ' ) if self . args . half else ( 32 , None )
classifier_config = classifier_config ,
convert_to = ' neuralnetwork ' if mlmodel else ' mlprogram ' )
bits , mode = ( 8 , ' kmeans ' ) if self . args . int8 else ( 16 , ' linear ' ) if self . args . half else ( 32 , None )
if bits < 32 :
if ' kmeans ' in mode :
check_requirements ( ' scikit-learn ' ) # scikit-learn package required for k-means quantization
ct_model = ct . models . neural_network . quantization_utils . quantize_weights ( ct_model , bits , mode )
if mlmodel :
ct_model = ct . models . neural_network . quantization_utils . quantize_weights ( ct_model , bits , mode )
else :
import coremltools . optimize . coreml as cto
op_config = cto . OpPalettizerConfig ( mode = mode , nbits = bits , weight_threshold = 512 )
config = cto . OptimizationConfig ( global_config = op_config )
ct_model = cto . palettize_weights ( ct_model , config = config )
if self . args . nms and self . model . task == ' detect ' :
ct_model = self . _pipeline_coreml ( ct_model )
if mlmodel :
import platform
# coremltools<=6.2 NMS export requires Python<3.11
check_version ( platform . python_version ( ) , ' <3.11 ' , name = ' Python ' , hard = True )
weights_dir = None
else :
ct_model . save ( str ( f ) ) # save otherwise weights_dir does not exist
weights_dir = str ( f / ' Data/com.apple.CoreML/weights ' )
ct_model = self . _pipeline_coreml ( ct_model , weights_dir = weights_dir )
m = self . metadata # metadata dict
ct_model . short_description = m . pop ( ' description ' )
@ -500,7 +522,14 @@ class Exporter:
ct_model . license = m . pop ( ' license ' )
ct_model . version = m . pop ( ' version ' )
ct_model . user_defined_metadata . update ( { k : str ( v ) for k , v in m . items ( ) } )
ct_model . save ( str ( f ) )
try :
ct_model . save ( str ( f ) ) # save *.mlpackage
except Exception as e :
LOGGER . warning (
f ' { prefix } WARNING ⚠️ { e } ), reverting to *.mlmodel export. '
f ' Known coremltools Python 3.11 and Windows bugs https://github.com/apple/coremltools/issues/1928. ' )
f = f . with_suffix ( ' .mlmodel ' )
ct_model . save ( str ( f ) )
return f , ct_model
@try_export
@ -546,7 +575,7 @@ class Exporter:
if self . args . dynamic :
shape = self . im . shape
if shape [ 0 ] < = 1 :
LOGGER . warning ( f ' { prefix } WARNING ⚠️ ' )
LOGGER . warning ( f " { prefix } WARNING ⚠️ ' dynamic=True ' model requires max batch size, i.e. ' batch=16 ' " )
profile = builder . create_optimization_profile ( )
for inp in inputs :
profile . set_shape ( inp . name , ( 1 , * shape [ 1 : ] ) , ( max ( 1 , shape [ 0 ] / / 2 ) , * shape [ 1 : ] ) , shape )
@ -805,7 +834,7 @@ class Exporter:
populator . populate ( )
tmp_file . unlink ( )
def _pipeline_coreml ( self , model , prefix = colorstr ( ' CoreML Pipeline: ' ) ) :
def _pipeline_coreml ( self , model , weights_dir = None , prefix = colorstr ( ' CoreML Pipeline: ' ) ) :
""" YOLOv8 CoreML pipeline. """
import coremltools as ct # noqa
@ -853,7 +882,7 @@ class Exporter:
# print(spec.description)
# Model from spec
model = ct . models . MLModel ( spec )
model = ct . models . MLModel ( spec , weights_dir = weights_dir )
# 3. Create NMS protobuf
nms_spec = ct . proto . Model_pb2 . Model ( )
@ -912,7 +941,7 @@ class Exporter:
' Confidence threshold ' : str ( nms . confidenceThreshold ) } )
# Save the model
model = ct . models . MLModel ( pipeline . spec )
model = ct . models . MLModel ( pipeline . spec , weights_dir = weights_dir )
model . input_description [ ' image ' ] = ' Input image '
model . input_description [ ' iouThreshold ' ] = f ' (optional) IOU threshold override (default: { nms . iouThreshold } ) '
model . input_description [ ' confidenceThreshold ' ] = \
Glenn Jocher
2 years ago
committed by
GitHub