export.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108
  1. # Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved.
  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. import os
  15. from paddle.utils import try_import
  16. __all__ = []
  17. def export(layer, path, input_spec=None, opset_version=9, **configs):
  18. """
  19. Export Layer to ONNX format, which can use for inference via onnxruntime or other backends.
  20. For more details, Please refer to `paddle2onnx <https://github.com/PaddlePaddle/paddle2onnx>`_ .
  21. Args:
  22. layer (Layer): The Layer to be exported.
  23. path (str): The path prefix to export model. The format is ``dirname/file_prefix`` or ``file_prefix`` ,
  24. and the exported ONNX file suffix is ``.onnx`` .
  25. input_spec (list[InputSpec|Tensor], optional): Describes the input of the exported model's forward
  26. method, which can be described by InputSpec or example Tensor. If None, all input variables of
  27. the original Layer's forward method would be the inputs of the exported ``ONNX`` model. Default: None.
  28. opset_version(int, optional): Opset version of exported ONNX model.
  29. Now, stable supported opset version include 9, 10, 11. Default: 9.
  30. **configs (dict, optional): Other export configuration options for compatibility. We do not
  31. recommend using these configurations, they may be removed in the future. If not necessary,
  32. DO NOT use them. Default None.
  33. The following options are currently supported:
  34. (1) output_spec (list[Tensor]): Selects the output targets of the exported model.
  35. By default, all return variables of original Layer's forward method are kept as the
  36. output of the exported model. If the provided ``output_spec`` list is not all output variables,
  37. the exported model will be pruned according to the given ``output_spec`` list.
  38. Returns:
  39. None
  40. Examples:
  41. .. code-block:: python
  42. >>> import paddle
  43. >>> class LinearNet(paddle.nn.Layer):
  44. ... def __init__(self):
  45. ... super().__init__()
  46. ... self._linear = paddle.nn.Linear(128, 10)
  47. ...
  48. ... def forward(self, x):
  49. ... return self._linear(x)
  50. ...
  51. >>> # Export model with 'InputSpec' to support dynamic input shape.
  52. >>> def export_linear_net():
  53. ... model = LinearNet()
  54. ... x_spec = paddle.static.InputSpec(shape=[None, 128], dtype='float32')
  55. ... paddle.onnx.export(model, 'linear_net', input_spec=[x_spec])
  56. ...
  57. >>> # doctest: +SKIP('raise ImportError')
  58. >>> export_linear_net()
  59. >>> class Logic(paddle.nn.Layer):
  60. ... def __init__(self):
  61. ... super().__init__()
  62. ...
  63. ... def forward(self, x, y, z):
  64. ... if z:
  65. ... return x
  66. ... else:
  67. ... return y
  68. ...
  69. >>> # Export model with 'Tensor' to support pruned model by set 'output_spec'.
  70. >>> def export_logic():
  71. ... model = Logic()
  72. ... x = paddle.to_tensor([1])
  73. ... y = paddle.to_tensor([2])
  74. ... # Static and run model.
  75. ... paddle.jit.to_static(model)
  76. ... out = model(x, y, z=True)
  77. ... paddle.onnx.export(model, 'pruned', input_spec=[x, y, z], output_spec=[out], input_names_after_prune=[x])
  78. ...
  79. >>> export_logic()
  80. """
  81. p2o = try_import('paddle2onnx')
  82. file_prefix = os.path.basename(path)
  83. if file_prefix == "":
  84. raise ValueError(
  85. "The input path MUST be format of dirname/file_prefix "
  86. "[dirname\\file_prefix in Windows system], but "
  87. f"the file_prefix is empty in received path: {path}"
  88. )
  89. save_file = path + '.onnx'
  90. p2o.dygraph2onnx(
  91. layer,
  92. save_file,
  93. input_spec=input_spec,
  94. opset_version=opset_version,
  95. **configs,
  96. )