(model,
im,
file,
dynamic,
tf_nms=False,
agnostic_nms=False,
topk_per_class=100,
topk_all=100,
iou_thres=0.45,
conf_thres=0.25,
keras=False,
prefix=colorstr('TensorFlow SavedModel:'))
| 302 | |
| 303 | @try_export |
| 304 | def export_saved_model(model, |
| 305 | im, |
| 306 | file, |
| 307 | dynamic, |
| 308 | tf_nms=False, |
| 309 | agnostic_nms=False, |
| 310 | topk_per_class=100, |
| 311 | topk_all=100, |
| 312 | iou_thres=0.45, |
| 313 | conf_thres=0.25, |
| 314 | keras=False, |
| 315 | prefix=colorstr('TensorFlow SavedModel:')): |
| 316 | # YOLOv5 TensorFlow SavedModel export |
| 317 | try: |
| 318 | import tensorflow as tf |
| 319 | except Exception: |
| 320 | check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}") |
| 321 | import tensorflow as tf |
| 322 | from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 |
| 323 | |
| 324 | from models.tf import TFModel |
| 325 | |
| 326 | LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') |
| 327 | f = str(file).replace('.pt', '_saved_model') |
| 328 | batch_size, ch, *imgsz = list(im.shape) # BCHW |
| 329 | |
| 330 | tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz) |
| 331 | im = tf.zeros((batch_size, *imgsz, ch)) # BHWC order for TensorFlow |
| 332 | _ = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) |
| 333 | inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if dynamic else batch_size) |
| 334 | outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) |
| 335 | keras_model = tf.keras.Model(inputs=inputs, outputs=outputs) |
| 336 | keras_model.trainable = False |
| 337 | keras_model.summary() |
| 338 | if keras: |
| 339 | keras_model.save(f, save_format='tf') |
| 340 | else: |
| 341 | spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype) |
| 342 | m = tf.function(lambda x: keras_model(x)) # full model |
| 343 | m = m.get_concrete_function(spec) |
| 344 | frozen_func = convert_variables_to_constants_v2(m) |
| 345 | tfm = tf.Module() |
| 346 | tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x), [spec]) |
| 347 | tfm.__call__(im) |
| 348 | tf.saved_model.save(tfm, |
| 349 | f, |
| 350 | options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) if check_version( |
| 351 | tf.__version__, '2.6') else tf.saved_model.SaveOptions()) |
| 352 | return f, keras_model |
| 353 | |
| 354 | |
| 355 | @try_export |
no test coverage detected