Skip to content

Commit

Permalink
Merge pull request #1376 from 0x3878f/recover-custom-ops
Browse files Browse the repository at this point in the history 
Support for customOp
  • Loading branch information
@risemeup1
risemeup1 authored Sep 6, 2024
2 parents 07acbcb + d3f0d0e commit f487f5c
Show file tree
Hide file tree
Showing 11 changed files with 230 additions and 31 deletions.
1 change: 1 addition & 0 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,6 +63,7 @@ paddle2onnx --model_dir saved_inference_model \
| --version | **[可选]** 查看 paddle2onnx 版本 |
| --external_filename | **[可选]** 当导出的 ONNX 模型大于 2G 时，需要设置 external data 的存储路径，推荐设置为：external_data |
| --export_fp16_model | **[可选]** 是否将导出的 ONNX 的模型转换为 FP16 格式，并用 ONNXRuntime-GPU 加速推理，默认为 False |
| --custom_ops | **[可选]** 将 Paddle OP 导出为 ONNX 的 Custom OP，例如：--custom_ops '{"paddle_op":"onnx_op"}，默认为 {} |


## 4.4 裁剪ONNX
Expand Down
3 changes: 2 additions & 1 deletion README_en.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,7 @@ The adjustable conversion parameters are listed in the following table:
| --version | **[Optional]** View paddle2onnx version |
| --external_filename | **[Optional]** When the exported ONNX model is larger than 2G, you need to set the storage path of external data, the recommended setting is: external_data |
| --export_fp16_model | **[Optional]** Whether to convert the exported ONNX model to FP16 format, and use ONNXRuntime-GPU to accelerate inference, the default is False |
| --custom_ops | **[Optional]** Export Paddle OP as ONNX's Custom OP, for example: --custom_ops '{"paddle_op":"onnx_op"}, default is {} |

## 4.4 Pruning ONNX

Expand All @@ -87,4 +88,4 @@ Provided under the [Apache-2.0 license](https://github.com/PaddlePaddle/paddle-o
# 7 Thank you for the donation

* Thanks to the PaddlePaddle team for providing server support for the CI infrastructure of Paddle2ONNX.
* Thanks to community users [chenwhql](https://github.com/chenwhql), [luotao1](https://github.com/luotao1), [goocody](https://github.com/goocody), [jeff41404](https://github.com/jeff41404), [jzhang553](https://github.com/jzhang533), [ZhengBicheng](https://github.com/ZhengBicheng) for donating a total of 10,000 RMB to the Paddle2ONNX PMC on March 28, 2024, for the development of Paddle2ONNX.
* Thanks to community users [chenwhql](https://github.com/chenwhql), [luotao1](https://github.com/luotao1), [goocody](https://github.com/goocody), [jeff41404](https://github.com/jeff41404), [jzhang553](https://github.com/jzhang533), [ZhengBicheng](https://github.com/ZhengBicheng) for donating a total of 10,000 RMB to the Paddle2ONNX PMC on March 28, 2024, for the development of Paddle2ONNX.
7 changes: 7 additions & 0 deletions paddle2onnx/command.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -99,6 +99,11 @@ def arg_parser():
type=ast.literal_eval,
default=False,
help="Whether export FP16 model for ORT-GPU, default False")
parser.add_argument(
"--custom_ops",
type=str,
default="{}",
help="Ops that needs to be converted to custom op, e.g --custom_ops '{\"paddle_op\":\"onnx_op\"}', default {}")
return parser


Expand Down Expand Up @@ -134,6 +139,7 @@ def main():
if base_path and not os.path.exists(base_path):
os.mkdir(base_path)
external_file = os.path.join(base_path, args.external_filename)
custom_ops_dict = eval(args.custom_ops)

calibration_file = args.save_calibration_file
paddle2onnx.export(
Expand All @@ -146,6 +152,7 @@ def main():
enable_onnx_checker=args.enable_onnx_checker,
enable_experimental_op=True,
enable_optimize=True,
custom_op_info=custom_ops_dict,
deploy_backend=args.deploy_backend,
calibration_file=calibration_file,
external_file=external_file,
Expand Down
26 changes: 12 additions & 14 deletions paddle2onnx/convert.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,25 +26,23 @@ def export(model_filename,
enable_onnx_checker=True,
enable_experimental_op=True,
enable_optimize=True,
custom_op_info=None,
deploy_backend="onnxruntime",
calibration_file="",
external_file="",
export_fp16_model=False):
deploy_backend = deploy_backend.lower()
# cpp2py_export.cc
onnx_model_str = c_p2o.export(
model_filename,
params_filename,
opset_version,
auto_upgrade_opset,
verbose,
enable_onnx_checker,
enable_experimental_op,
enable_optimize,
deploy_backend,
calibration_file,
external_file,
export_fp16_model)
if custom_op_info is None:
onnx_model_str = c_p2o.export(
model_filename, params_filename, opset_version, auto_upgrade_opset, verbose,
enable_onnx_checker, enable_experimental_op, enable_optimize, {},
deploy_backend, calibration_file, external_file, export_fp16_model)
else:
onnx_model_str = c_p2o.export(
model_filename, params_filename, opset_version, auto_upgrade_opset, verbose,
enable_onnx_checker, enable_experimental_op, enable_optimize,
custom_op_info, deploy_backend, calibration_file, external_file,
export_fp16_model)
if save_file is not None:
with open(save_file, "wb") as f:
f.write(onnx_model_str)
Expand Down
29 changes: 29 additions & 0 deletions paddle2onnx/converter.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,8 @@ namespace paddle2onnx
bool enable_onnx_checker,
bool enable_experimental_op,
bool enable_optimize,
CustomOp* ops,
int op_count,
const char *deploy_backend,
char **calibration_cache,
int *calibration_size,
Expand All @@ -53,6 +55,19 @@ namespace paddle2onnx
}
paddle2onnx::ModelExporter me;

// Add custom operator information
if (ops != nullptr && op_count > 0) {
for (int i = 0; i < op_count; ++i) {
std::string op_name(ops[i].op_name, strlen(ops[i].op_name));
std::string export_op_name(ops[i].export_op_name,
strlen(ops[i].export_op_name));
if (export_op_name == "paddle2onnx_null") {
export_op_name = op_name;
}
me.custom_ops[op_name] = export_op_name;
}
}

// Add disabled fp16 op information
std::vector<std::string> disable_op_types;
if (disable_fp16_op_types != nullptr && disable_fp16_op_types_count > 0)
Expand Down Expand Up @@ -107,6 +122,8 @@ namespace paddle2onnx
bool enable_onnx_checker,
bool enable_experimental_op,
bool enable_optimize,
CustomOp* ops,
int op_count,
const char *deploy_backend,
char **calibration_cache,
int *calibration_size,
Expand All @@ -124,6 +141,18 @@ namespace paddle2onnx
return false;
}
paddle2onnx::ModelExporter me;
// Add custom operator information
if (ops != nullptr && op_count > 0) {
for (int i = 0; i < op_count; ++i) {
std::string op_name(ops[i].op_name, strlen(ops[i].op_name));
std::string export_op_name(ops[i].export_op_name,
strlen(ops[i].export_op_name));
if (export_op_name == "paddle2onnx_null") {
export_op_name = op_name;
}
me.custom_ops[op_name] = export_op_name;
}
}

// Add disabled fp16 op information
std::vector<std::string> disable_op_types;
Expand Down
4 changes: 4 additions & 0 deletions paddle2onnx/converter.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,8 @@ namespace paddle2onnx
bool enable_onnx_checker = true,
bool enable_experimental_op = false,
bool enable_optimize = true,
CustomOp* ops = nullptr,
int op_count = 0,
const char *deploy_backend = "onnxruntime",
char **calibration_cache = nullptr,
int *calibration_size = 0,
Expand All @@ -65,6 +67,8 @@ namespace paddle2onnx
bool enable_onnx_checker = true,
bool enable_experimental_op = false,
bool enable_optimize = true,
CustomOp* ops = nullptr,
int op_count = 0,
const char *deploy_backend = "onnxruntime",
char **calibration_cache = nullptr,
int *calibration_size = 0,
Expand Down
66 changes: 54 additions & 12 deletions paddle2onnx/cpp2py_export.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,39 +29,81 @@ PYBIND11_MODULE(paddle2onnx_cpp2py_export, m) {
m.doc() = "Paddle2ONNX: export PaddlePaddle to ONNX";
// converter.cc
m.def("export", [](const std::string& model_filename,
const std::string& params_filename,
const std::string& params_filename,
int opset_version = 7,
bool auto_upgrade_opset = true,
bool auto_upgrade_opset = true,
bool verbose = true,
bool enable_onnx_checker = true,
bool enable_experimental_op = true,
bool enable_optimize = true,
const CustomOpInfo& info = CustomOpInfo(),
const std::string& deploy_backend = "onnxruntime",
const std::string& calibration_file = "",
const std::string& external_file = "",
const bool& export_fp16_model = false) {
P2OLogger(verbose) << "Start to parse PaddlePaddle model..." << std::endl;
P2OLogger(verbose) << "Model file path: " << model_filename << std::endl;
P2OLogger(verbose) << "Parameters file path: " << params_filename << std::endl;
if (info.size() == 0) {
char* out = nullptr;
int size = 0;
char* calibration_cache = nullptr;
int cache_size = 0;
bool save_external;
if (!Export(model_filename.c_str(), params_filename.c_str(), &out, &size,
opset_version, auto_upgrade_opset, verbose,
enable_onnx_checker, enable_experimental_op, enable_optimize,
nullptr, 0, deploy_backend.c_str(), &calibration_cache,
&cache_size, external_file.c_str(), &save_external,
export_fp16_model)) {
P2OLogger(verbose) << "Paddle model convert failed." << std::endl;
return pybind11::bytes("");
}
if (cache_size) {
std::string calibration_cache_str(calibration_cache,
calibration_cache + cache_size);
std::ofstream cache_file;
cache_file.open(calibration_file, std::ios::out);
cache_file << calibration_cache_str;
delete calibration_cache;
calibration_cache = nullptr;
P2OLogger(verbose) << "TensorRT calibration cache path: "
<< calibration_file << std::endl;
}
std::string onnx_proto(out, out + size);
delete out;
out = nullptr;
return pybind11::bytes(onnx_proto);
}
std::vector<CustomOp> ops;
ops.resize(info.size());
int index = 0;
for (auto& item : info) {
strcpy(ops[index].op_name, item.first.c_str());
strcpy(ops[index].export_op_name, item.second.c_str());
index += 1;
}
char* out = nullptr;
int size = 0;
char* calibration_cache = nullptr;
int cache_size = 0;
bool save_external;
if (!Export(model_filename.c_str(),
params_filename.c_str(),
&out,
if (!Export(model_filename.c_str(),
params_filename.c_str(),
&out,
&size,
opset_version,
auto_upgrade_opset,
opset_version,
auto_upgrade_opset,
verbose,
enable_onnx_checker,
enable_experimental_op,
enable_onnx_checker,
enable_experimental_op,
enable_optimize,
deploy_backend.c_str(),
ops.data(),
info.size(),
deploy_backend.c_str(),
&calibration_cache,
&cache_size,
external_file.c_str(),
&cache_size,
external_file.c_str(),
&save_external,
export_fp16_model)) {
P2OLogger(verbose) << "Paddle model convert failed." << std::endl;
Expand Down
Loading

0 comments on commit f487f5c

Please sign in to comment.