diff --git a/pytorch2keras/converter.py b/pytorch2keras/converter.py index 3767099..6f27b63 100644 --- a/pytorch2keras/converter.py +++ b/pytorch2keras/converter.py @@ -5,7 +5,6 @@ from onnx2keras import onnx_to_keras import torch import onnx -from onnx import optimizer import io import logging @@ -13,7 +12,7 @@ def pytorch_to_keras( model, args, input_shapes=None, change_ordering=False, verbose=False, name_policy=None, - use_optimizer=False, do_constant_folding=False + do_constant_folding=False ): """ By given PyTorch model convert layers with ONNX. @@ -69,17 +68,5 @@ def pytorch_to_keras( stream.seek(0) onnx_model = onnx.load(stream) - if use_optimizer: - if use_optimizer is True: - optimizer2run = optimizer.get_available_passes() - else: - use_optimizer = set(use_optimizer) - optimizer2run = [x for x in optimizer.get_available_passes() if x in use_optimizer] - logger.info("Running optimizer:\n%s", "\n".join(optimizer2run)) - onnx_model = optimizer.optimize(onnx_model, optimizer2run) - - k_model = onnx_to_keras(onnx_model=onnx_model, input_names=input_names, - input_shapes=input_shapes, name_policy=name_policy, - verbose=verbose, change_ordering=change_ordering) return k_model