Skip to content

Commit

Permalink
Add --keras argument for TF exports (ultralytics#7921)
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher authored and Clay Januhowski committed Sep 8, 2022
1 parent aca424b commit 94787f4
Showing 1 changed file with 4 additions and 1 deletion.
5 changes: 4 additions & 1 deletion export.py
Original file line number Diff line number Diff line change
Expand Up @@ -574,6 +574,7 @@ def run(
half=False, # FP16 half-precision export
inplace=False, # set YOLOv5 Detect() inplace=True
train=False, # model.train() mode
keras=False, # use Keras
optimize=False, # TorchScript: optimize for mobile
int8=False, # CoreML/TF INT8 quantization
dynamic=False, # ONNX/TF: dynamic axes
Expand Down Expand Up @@ -656,8 +657,9 @@ def run(
agnostic_nms=agnostic_nms or tfjs,
topk_per_class=topk_per_class,
topk_all=topk_all,
iou_thres=iou_thres,
conf_thres=conf_thres,
iou_thres=iou_thres) # keras model
keras=keras)
if pb or tfjs: # pb prerequisite to tfjs
f[6] = export_pb(model, file)
if tflite or edgetpu:
Expand Down Expand Up @@ -689,6 +691,7 @@ def parse_opt():
parser.add_argument('--half', action='store_true', help='FP16 half-precision export')
parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
parser.add_argument('--train', action='store_true', help='model.train() mode')
parser.add_argument('--keras', action='store_true', help='TF: use Keras')
parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile')
parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')
Expand Down

0 comments on commit 94787f4

Please sign in to comment.