expose get_all_supported_ops in python (#1374)

Signed-off-by: Zhang Jun <jzhang533@gmail.com>
PaddlePaddle · Sep 7, 2024 · 0a5915e · 0a5915e
1 parent f487f5c
commit 0a5915e
Show file tree

Hide file tree

Showing 4 changed files with 21 additions and 15 deletions.
diff --git a/docs/zh/Paddle2ONNX_Development_Guide.md b/docs/zh/Paddle2ONNX_Development_Guide.md
@@ -53,6 +53,14 @@ Paddle2ONNX 开发的主要步骤为：
 
 通过 Netron 的可视化，可以看到， **ATTRIBUTES** 中的参数 **axis** 和 **shifts** 与核心文档中的输入参数一一对应。
 
+> [!NOTE]
+> 获取当前所有支持的 OP 列表的方法：
+>
+> ```python
+> import paddle2onnx
+> paddle2onnx.get_all_supported_operators()
+> ```
+
 ### 3.3 查阅 ONNX API 文档
 
 掌握 Paddle OP 的原理和使用方式后，查阅 [ONNX Operators Docs](https://onnx.ai/onnx/operators/index.html) 找到对应的实现，若 ONNX OP 和 Paddle OP 没有一对一的实现，则需要根据 Paddle OP 的原理使用多个 ONNX OP 组合实现。 

diff --git a/paddle2onnx/cpp2py_export.cc b/paddle2onnx/cpp2py_export.cc
@@ -21,7 +21,7 @@
 #include "paddle2onnx/converter.h"
 #include "paddle2onnx/mapper/exporter.h"
 #include "paddle2onnx/optimizer/paddle2onnx_optimizer.h"
-
+#include "paddle2onnx/mapper/register_mapper.h"
 namespace paddle2onnx {
 
 typedef std::map<std::string, std::string> CustomOpInfo;
@@ -137,5 +137,10 @@ PYBIND11_MODULE(paddle2onnx_cpp2py_export, m) {
     ONNX_NAMESPACE::optimization::Paddle2ONNXFP32ToFP16(fp32_model_path,
                                                      fp16_model_path);
   });
+  m.def("get_all_supported_operators", []() {
+
+    auto operators = MapperHelper::Get()->GetAllOps();
+    return operators;
+  });
 }
 }  // namespace paddle2onnx
diff --git a/paddle2onnx/mapper/register_mapper.h b/paddle2onnx/mapper/register_mapper.h
@@ -63,20 +63,12 @@ class MapperHelper {
     return helper;
   }
 
-  int64_t GetAllOps(const std::string& file_path) {
-    std::ofstream outfile(file_path);
-    if (!outfile) {
-      std::cerr << "Failed to open file: " << file_path << std::endl;
-      return mappers.size();
-    }
+  std::vector<std::string> GetAllOps() {
+    std::vector<std::string> operators;
     for (auto iter = mappers.begin(); iter != mappers.end(); iter++) {
-      outfile << iter->first << std::endl;
+      operators.push_back(iter->first);
     }
-    outfile << "Total OPs: " << mappers.size() << std::endl;
-    std::cout << " [ * Paddle2ONNX * ] All Registered OPs saved in "
-              << file_path << std::endl;
-    outfile.close();
-    return mappers.size();
+    return operators;
   }
 
   bool IsRegistered(const std::string& op_name) {

diff --git a/paddle2onnx/utils.py b/paddle2onnx/utils.py
@@ -15,11 +15,12 @@
 from __future__ import absolute_import
 
 import importlib
-import collections
 import time
-import os
 import sys
+import paddle2onnx.paddle2onnx_cpp2py_export as c_p2o
 
+def get_all_supported_operators():
+    return c_p2o.get_all_supported_operators()
 
 def try_import(module_name):
     """Try importing a module, with an informative error message on failure."""