convertor.py 1.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142
  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from __future__ import absolute_import
  15. import paddle.fluid as fluid
  16. import os
  17. import sys
  18. import paddlex as pdx
  19. import paddlex.utils.logging as logging
  20. __all__ = ['export_onnx']
  21. def export_onnx(model_dir, save_dir, fixed_input_shape):
  22. assert len(fixed_input_shape) == 2, "len of fixed input shape must == 2"
  23. model = pdx.load_model(model_dir, fixed_input_shape)
  24. model_name = os.path.basename(model_dir.strip('/')).split('/')[-1]
  25. export_onnx_model(model, save_dir)
  26. def export_onnx_model(model, save_dir):
  27. try:
  28. import x2paddle
  29. if x2paddle.__version__ < '0.7.4':
  30. logging.error("You need to upgrade x2paddle >= 0.7.4")
  31. except:
  32. logging.error(
  33. "You need to install x2paddle first, pip install x2paddle>=0.7.4")
  34. from x2paddle.op_mapper.paddle_op_mapper import PaddleOpMapper
  35. mapper = PaddleOpMapper()
  36. mapper.convert(model.test_prog, save_dir)