Merge pull request #3 from ddding-z/master

ReTree Random Forest Implementation

Author: yangxinyu7427

Diff for: .gitmodules

@@ -3,4 +3,4 @@
 	url = https://github.com/protocolbuffers/protobuf.git
 [submodule "third_party/onnx"]
 	path = third_party/onnx
-	url = git@github.com:lovelynewlife/onnx.git
+	url = https://github.com/lovelynewlife/onnx.git

Diff for: CMakeLists.txt

@@ -1,4 +1,4 @@
-cmake_minimum_required(VERSION 3.22)
+cmake_minimum_required(VERSION 3.16)
 
 # For std::filesystem
 # Must be a cache variable and be set before project()

Diff for: build.sh

@@ -7,4 +7,4 @@ rm -rf build
 cmake -B build -DCMAKE_BUILD_TYPE=Debug -DCMAKE_EXPORT_COMPILE_COMMANDS=ON
 
 # 构建项目，Debug 模式，先清理再并行编译
-cmake --build build --config Debug --clean-first --parallel 10
+cmake --build build --config Debug --clean-first --parallel 80

Diff for: onnxoptimizer/optimize_c_api/optimize_c_api.cpp

@@ -54,25 +54,24 @@ void add_prefix_on_model(std::string& changed_model_path, std::string& output_mo
     onnx::optimization::add_prefix_on_model(changed_model_path, output_model_path, prefix);
 }
 
-/// @brief 
-/// @param input_model_path 
-/// @param comparison_operator 1: ==, 2: <, 3: <=, 4: >, 5: >= 
-/// @param threshold 
-/// @param features
-/// @return optimized-model path
 std::string optimize_on_decision_tree_predicate(std::string& input_model_path, uint8_t comparison_operator,
-                float threshold) {
+                float threshold, int threads_count) {
     std::string mp1 = onnx::optimization::DTConvertRule::match(input_model_path);
-	std::string mp2 = onnx::optimization::DTPruneRule::match(mp1, comparison_operator, threshold);
-	return onnx::optimization::DTMergeRule::match(mp2);
+	std::string mp2 = onnx::optimization::DTPruneRule::match(mp1, comparison_operator, threshold, threads_count);
+	return onnx::optimization::DTMergeRule::match(mp2, threads_count);
+}
 
-	// std::string mp1 = onnx::optimization::DTPruneRule::match(input_model_path, comparison_operator, threshold);
-	// return onnx::optimization::DTMergeRule::match(mp1);
-	// return onnx::optimization::DTMergeRule::match(mp3);
-    // return onnx::optimization::DTConvertRule::match(input_model_path);
-    // return onnx::optimization::DTPruneRule::match(input_model_path, comparison_operator, threshold);
+std::string optimize_on_decision_tree_predicate_convert(std::string& input_model_path){
+    return onnx::optimization::DTConvertRule::match(input_model_path);
+}
 
-    // return onnx::optimization::DTConvertRule::match(input_model_path);
+std::string optimize_on_decision_tree_predicate_prune(std::string& input_model_path, uint8_t comparison_operator, float threshold, int threads_count){
+    return onnx::optimization::DTPruneRule::match(input_model_path, comparison_operator, threshold, threads_count);
 }
 
+std::string optimize_on_decision_tree_predicate_merge(std::string& input_model_path, int threads_count){
+    return onnx::optimization::DTMergeRule::match(input_model_path, threads_count);
+}
+
+
 // -----------------------

Diff for: onnxoptimizer/optimize_c_api/optimize_c_api.h

@@ -34,8 +34,23 @@ void change_models(std::string& changed_model_path,std::string& output_model_pat
                 const std::vector<std::string>& output_name);
 void add_prefix_on_model(std::string& changed_model_path, std::string& output_model_path, std::string& prefix);
 
-//--------------------------
-std::string optimize_on_decision_tree_predicate(std::string& input_model_path, uint8_t comparison_operator, float threshold);
+/**
+ * @brief retree optimization 
+ *        1. convert
+ *        2. prune
+ *        3. merge
+ * @param input_model_path
+ * @param comparison_operator 1: ==, 2: <, 3: <=, 4: >, 5: >= 
+ * @param threshold 
+ * @param threads_count default=1 
+ * @return optimized model path
+ */
+std::string optimize_on_decision_tree_predicate(std::string& input_model_path, uint8_t comparison_operator, float threshold, int threads_count = 1);
+
+std::string optimize_on_decision_tree_predicate_convert(std::string& input_model_path);
+std::string optimize_on_decision_tree_predicate_prune(std::string& input_model_path, uint8_t comparison_operator, float threshold, int threads_count = 1);
+std::string optimize_on_decision_tree_predicate_merge(std::string& input_model_path, int threads_count = 1);
+
 
 #endif  // ONNX_OPTIMIZER_OPTIMIZE_C_API_H