@@ -434,42 +434,45 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path'
434
434
LOGGER .info (f"\n { colorstr ('PyTorch:' )} starting from { file } ({ file_size (file ):.1f} MB)" )
435
435
436
436
# Exports
437
+ f = ['' ] * 10 # exported filenames
437
438
if 'torchscript' in include :
438
- f = export_torchscript (model , im , file , optimize )
439
+ f [ 0 ] = export_torchscript (model , im , file , optimize )
439
440
if 'engine' in include : # TensorRT required before ONNX
440
- f = export_engine (model , im , file , train , half , simplify , workspace , verbose )
441
+ f [ 1 ] = export_engine (model , im , file , train , half , simplify , workspace , verbose )
441
442
if ('onnx' in include ) or ('openvino' in include ): # OpenVINO requires ONNX
442
- f = export_onnx (model , im , file , opset , train , dynamic , simplify )
443
+ f [ 2 ] = export_onnx (model , im , file , opset , train , dynamic , simplify )
443
444
if 'openvino' in include :
444
- f = export_openvino (model , im , file )
445
+ f [ 3 ] = export_openvino (model , im , file )
445
446
if 'coreml' in include :
446
- _ , f = export_coreml (model , im , file )
447
+ _ , f [ 4 ] = export_coreml (model , im , file )
447
448
448
449
# TensorFlow Exports
449
450
if any (tf_exports ):
450
451
pb , tflite , edgetpu , tfjs = tf_exports [1 :]
451
452
if int8 or edgetpu : # TFLite --int8 bug https://github.com/ultralytics/yolov5/issues/5707
452
453
check_requirements (('flatbuffers==1.12' ,)) # required before `import tensorflow`
453
454
assert not (tflite and tfjs ), 'TFLite and TF.js models must be exported separately, please pass only one type.'
454
- model , f = export_saved_model (model , im , file , dynamic , tf_nms = nms or agnostic_nms or tfjs ,
455
- agnostic_nms = agnostic_nms or tfjs , topk_per_class = topk_per_class ,
456
- topk_all = topk_all , conf_thres = conf_thres , iou_thres = iou_thres ) # keras model
455
+ model , f [ 5 ] = export_saved_model (model , im , file , dynamic , tf_nms = nms or agnostic_nms or tfjs ,
456
+ agnostic_nms = agnostic_nms or tfjs , topk_per_class = topk_per_class ,
457
+ topk_all = topk_all , conf_thres = conf_thres , iou_thres = iou_thres ) # keras model
457
458
if pb or tfjs : # pb prerequisite to tfjs
458
- f = export_pb (model , im , file )
459
+ f [ 6 ] = export_pb (model , im , file )
459
460
if tflite or edgetpu :
460
- f = export_tflite (model , im , file , int8 = int8 or edgetpu , data = data , ncalib = 100 )
461
+ f [ 7 ] = export_tflite (model , im , file , int8 = int8 or edgetpu , data = data , ncalib = 100 )
461
462
if edgetpu :
462
- f = export_edgetpu (model , im , file )
463
+ f [ 8 ] = export_edgetpu (model , im , file )
463
464
if tfjs :
464
- f = export_tfjs (model , im , file )
465
+ f [ 9 ] = export_tfjs (model , im , file )
465
466
466
467
# Finish
468
+ f = [str (x ) for x in f if x ] # filter out '' and None
467
469
LOGGER .info (f'\n Export complete ({ time .time () - t :.2f} s)'
468
470
f"\n Results saved to { colorstr ('bold' , file .parent .resolve ())} "
469
471
f"\n Visualize with https://netron.app"
470
- f"\n Detect with `python detect.py --weights { f } `"
471
- f" or `model = torch.hub.load('ultralytics/yolov5', 'custom', '{ f } ')"
472
- f"\n Validate with `python val.py --weights { f } `" )
472
+ f"\n Detect with `python detect.py --weights { f [- 1 ]} `"
473
+ f" or `model = torch.hub.load('ultralytics/yolov5', 'custom', '{ f [- 1 ]} ')"
474
+ f"\n Validate with `python val.py --weights { f [- 1 ]} `" )
475
+ return f # return list of exported files/dirs
473
476
474
477
475
478
def parse_opt ():
0 commit comments