glenn-jocher commited on
Commit
547c89b
·
unverified ·
1 Parent(s): 15057d2

Add `--keras` argument for TF exports (#7921)

Browse files

Resolves https://github.com/ultralytics/yolov5/issues/7911#issuecomment-1133671255

Files changed (1) hide show
  1. export.py +4 -1
export.py CHANGED
@@ -455,6 +455,7 @@ def run(
455
  half=False, # FP16 half-precision export
456
  inplace=False, # set YOLOv5 Detect() inplace=True
457
  train=False, # model.train() mode
 
458
  optimize=False, # TorchScript: optimize for mobile
459
  int8=False, # CoreML/TF INT8 quantization
460
  dynamic=False, # ONNX/TF: dynamic axes
@@ -536,8 +537,9 @@ def run(
536
  agnostic_nms=agnostic_nms or tfjs,
537
  topk_per_class=topk_per_class,
538
  topk_all=topk_all,
 
539
  conf_thres=conf_thres,
540
- iou_thres=iou_thres) # keras model
541
  if pb or tfjs: # pb prerequisite to tfjs
542
  f[6] = export_pb(model, file)
543
  if tflite or edgetpu:
@@ -569,6 +571,7 @@ def parse_opt():
569
  parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
570
  parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
571
  parser.add_argument('--train', action='store_true', help='model.train() mode')
 
572
  parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
573
  parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
574
  parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
 
455
  half=False, # FP16 half-precision export
456
  inplace=False, # set YOLOv5 Detect() inplace=True
457
  train=False, # model.train() mode
458
+ keras=False, # use Keras
459
  optimize=False, # TorchScript: optimize for mobile
460
  int8=False, # CoreML/TF INT8 quantization
461
  dynamic=False, # ONNX/TF: dynamic axes
 
537
  agnostic_nms=agnostic_nms or tfjs,
538
  topk_per_class=topk_per_class,
539
  topk_all=topk_all,
540
+ iou_thres=iou_thres,
541
  conf_thres=conf_thres,
542
+ keras=keras)
543
  if pb or tfjs: # pb prerequisite to tfjs
544
  f[6] = export_pb(model, file)
545
  if tflite or edgetpu:
 
571
  parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
572
  parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
573
  parser.add_argument('--train', action='store_true', help='model.train() mode')
574
+ parser.add_argument('--keras', action='store_true', help='TF: use Keras')
575
  parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
576
  parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
577
  parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')