Skip to content

Commit 20f2803

Browse files
committed
Allow to select optimizers in from_keras conversion function
Signed-off-by: Sagar Shelke <[email protected]>
1 parent 07e9e0f commit 20f2803

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

tf2onnx/convert.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -399,7 +399,7 @@ def _from_keras_tf1(model, opset=None, custom_ops=None, custom_op_handlers=None,
399399

400400
def from_keras(model, input_signature=None, opset=None, custom_ops=None, custom_op_handlers=None,
401401
custom_rewriter=None, inputs_as_nchw=None, extra_opset=None, shape_override=None,
402-
target=None, large_model=False, output_path=None):
402+
target=None, large_model=False, output_path=None, optimizers=None):
403403
"""Returns a ONNX model_proto for a tf.keras model.
404404
405405
Args:
@@ -417,6 +417,7 @@ def from_keras(model, input_signature=None, opset=None, custom_ops=None, custom_
417417
inputs_as_nchw: transpose inputs in list from nchw to nhwc
418418
large_model: use the ONNX external tensor storage format
419419
output_path: save model to output_path
420+
optimizers: list (subset) of tf2onnx optimizers if applying all optimizers is not desired.
420421
421422
Returns:
422423
An ONNX model_proto and an external_tensor_storage dict.
@@ -489,6 +490,7 @@ def wrap_call(*args, training=False, **kwargs):
489490
opset=opset,
490491
custom_ops=custom_ops,
491492
custom_op_handlers=custom_op_handlers,
493+
optimizers=optimizers,
492494
custom_rewriter=custom_rewriter,
493495
extra_opset=extra_opset,
494496
shape_override=shape_override,

0 commit comments

Comments
 (0)