-
Notifications
You must be signed in to change notification settings - Fork 3k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add ONNX Export #2061
Add ONNX Export #2061
Conversation
paddlenlp/trainer/trainer_base.py
Outdated
|
||
Args: | ||
input_spec (paddle.static.InputSpec, optional): InputSpec describes the signature information of the model input, | ||
such as shape , dtype , name. Defaults to None. | ||
load_best_model (bool, optional): Load best model. Defaults to False. | ||
output_dir (Optional[str], optional): Output dir to save the exported model. Defaults to None. | ||
export_model_format (Optional[str], optional): Export model format. Defaults to Paddle. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
注释写清楚有哪些可选项
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
save_path = os.path.join(output_dir, "onnx", "model") | ||
logger.info("Exporting ONNX model to %s" % save_path) | ||
paddle.onnx.export(model, save_path, input_spec=input_spec) | ||
logger.info("ONNX model exported.") |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
处理else的情况,给出明确的提示
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
paddlenlp/trainer/trainer_base.py
Outdated
output_dir: Optional[str]=None): | ||
""" Export paddle inference model. | ||
output_dir: Optional[str]=None, | ||
export_model_format: Optional[str]="Paddle"): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
export_model_format
建议小写,内部处理的时候,再统一转一下小写。
export_model_format = export_model_format.lower()
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
LGTM
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
似乎没看见在requirements中新增paddle2onnx依赖
paddlenlp/trainer/trainer_base.py
Outdated
output_dir: Optional[str]=None): | ||
""" Export paddle inference model. | ||
output_dir: Optional[str]=None, | ||
export_model_format: Optional[str]="paddle"): |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
export_model_format -> model_format
因为函数已经表示当作export,参数变量更多是修饰这个动作即可。
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Done.并在requirements中新增paddle2onnx依赖
PR types
New features
PR changes
Others
Description
Add ONNX model export option