zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

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

Resolves https://github.com/ultralytics/yolov5/issues/7911#issuecomment-1133671255
This commit is contained in:
Glenn Jocher 2022-05-22 13:33:36 +02:00 committed by GitHub
parent 15057d2a34
commit 547c89b3a0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
export.py
View File

@ -455,6 +455,7 @@ def run(
half=False, # FP16 half-precision export half=False, # FP16 half-precision export
inplace=False, # set YOLOv5 Detect() inplace=True inplace=False, # set YOLOv5 Detect() inplace=True
train=False, # model.train() mode train=False, # model.train() mode
keras=False, # use Keras
optimize=False, # TorchScript: optimize for mobile optimize=False, # TorchScript: optimize for mobile
int8=False, # CoreML/TF INT8 quantization int8=False, # CoreML/TF INT8 quantization
dynamic=False, # ONNX/TF: dynamic axes dynamic=False, # ONNX/TF: dynamic axes
@ -536,8 +537,9 @@ def run(
agnostic_nms=agnostic_nms or tfjs, agnostic_nms=agnostic_nms or tfjs,
topk_per_class=topk_per_class, topk_per_class=topk_per_class,
topk_all=topk_all, topk_all=topk_all,
iou_thres=iou_thres,
conf_thres=conf_thres, conf_thres=conf_thres,
iou_thres=iou_thres) # keras model keras=keras)
if pb or tfjs: # pb prerequisite to tfjs if pb or tfjs: # pb prerequisite to tfjs
f[6] = export_pb(model, file) f[6] = export_pb(model, file)
if tflite or edgetpu: if tflite or edgetpu:
@ -569,6 +571,7 @@ def parse_opt():
parser.add_argument('--half', action='store_true', help='FP16 half-precision export') 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('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True')
parser.add_argument('--train', action='store_true', help='model.train() mode') 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('--optimize', action='store_true', help='TorchScript: optimize for mobile')
parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization') parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization')
parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes') parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes')