diff --git a/onnxruntime/core/providers/cpu/ml/treeregressor.cc b/onnxruntime/core/providers/cpu/ml/treeregressor.cc index 2bb0a3db3c..0c96ffaaa0 100644 --- a/onnxruntime/core/providers/cpu/ml/treeregressor.cc +++ b/onnxruntime/core/providers/cpu/ml/treeregressor.cc @@ -147,7 +147,7 @@ TreeEnsembleRegressor::TreeEnsembleRegressor(const OpKernelInfo& info) } template -common::Status TreeEnsembleRegressor::ProcessTreeNode(std::unordered_map& classes, int64_t treeindex, const T* Xdata, int64_t feature_base) const { +common::Status TreeEnsembleRegressor::ProcessTreeNode(std::unordered_map < int64_t, std::tuple>& classes, int64_t treeindex, const T* Xdata, int64_t feature_base) const { //walk down tree to the leaf ::onnxruntime::ml::NODE_MODE mode = static_cast<::onnxruntime::ml::NODE_MODE>(nodes_modes_[treeindex]); int64_t loopcount = 0; @@ -197,9 +197,13 @@ common::Status TreeEnsembleRegressor::ProcessTreeNode(std::unordered_map(leafnode_data_[index]); auto it_classes = classes.find(classid); if (it_classes != classes.end()) { - it_classes->second += weight; + auto& tuple = it_classes->second; + std::get<0>(tuple) += weight; + if (weight < std::get<1>(tuple)) std::get<1>(tuple) = weight; + if (weight > std::get<2>(tuple)) std::get<2>(tuple) = weight; } else { - auto p1 = std::make_pair(classid, weight); + std::tuple tuple = std::make_tuple(weight, weight, weight); + auto p1 = std::make_pair(classid, tuple); classes.insert(p1); } index++; @@ -232,7 +236,7 @@ common::Status TreeEnsembleRegressor::Compute(OpKernelContext* context) const for (int64_t i = 0; i < N; i++) //for each class { int64_t current_weight_0 = i * stride; - std::unordered_map scores; + std::unordered_map> scores; // sum, min, max //for each tree for (size_t j = 0; j < roots_.size(); j++) { //walk each tree from its root @@ -246,13 +250,13 @@ common::Status TreeEnsembleRegressor::Compute(OpKernelContext* context) const float val = base_values_.size() == (size_t)n_targets_ ? base_values_[j] : 0.f; if (it_scores != scores.end()) { if (aggregate_function_ == ::onnxruntime::ml::AGGREGATE_FUNCTION::AVERAGE) { - val += scores[j] / roots_.size(); + val += std::get<0>(scores[j]) / roots_.size(); //first element of tuple is already a sum } else if (aggregate_function_ == ::onnxruntime::ml::AGGREGATE_FUNCTION::SUM) { - val += scores[j]; + val += std::get<0>(scores[j]); } else if (aggregate_function_ == ::onnxruntime::ml::AGGREGATE_FUNCTION::MIN) { - if (scores[j] < val) val = scores[j]; + val += std::get<1>(scores[j]); // second element of tuple is min } else if (aggregate_function_ == ::onnxruntime::ml::AGGREGATE_FUNCTION::MAX) { - if (scores[j] > val) val = scores[j]; + val += std::get<2>(scores[j]); // third element of tuple is max } } outputs.push_back(val); diff --git a/onnxruntime/core/providers/cpu/ml/treeregressor.h b/onnxruntime/core/providers/cpu/ml/treeregressor.h index 2216e5c74b..2af61fe393 100644 --- a/onnxruntime/core/providers/cpu/ml/treeregressor.h +++ b/onnxruntime/core/providers/cpu/ml/treeregressor.h @@ -15,7 +15,7 @@ class TreeEnsembleRegressor final : public OpKernel { common::Status Compute(OpKernelContext* context) const override; private: - common::Status ProcessTreeNode(std::unordered_map& classes, int64_t treeindex, const T* Xdata, int64_t feature_base) const; + common::Status ProcessTreeNode(std::unordered_map < int64_t, std::tuple>& classes, int64_t treeindex, const T* Xdata, int64_t feature_base) const; std::vector nodes_treeids_; std::vector nodes_nodeids_; diff --git a/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc b/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc index 335c17b856..31cfcc7cad 100644 --- a/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc +++ b/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc @@ -7,7 +7,8 @@ namespace onnxruntime { namespace test { -TEST(MLOpTest, TreeRegressorMultiTarget) { +void GenTreeAndRunTest(const std::vector& X, const std::vector& base_values, const std::vector& results, const std::string& aggFunction) +{ OpTester test("TreeEnsembleRegressor", 1, onnxruntime::kMLDomain); //tree @@ -25,9 +26,78 @@ TEST(MLOpTest, TreeRegressorMultiTarget) { std::vector target_weights = {1.5f, 27.5f, 2.25f, 20.75f, 2.f, 23.f, 3.f, 14.f, 0.f, 41.f, 1.83333333f, 24.5f, 0.f, 41.f, 2.75f, 16.25f, 2.f, 23.f, 3.f, 14.f, 2.66666667f, 17.f, 2.f, 23.f, 3.f, 14.f}; std::vector classes = {0, 1}; - //test data + //add attributes + test.AddAttribute("nodes_truenodeids", lefts); + test.AddAttribute("nodes_falsenodeids", rights); + test.AddAttribute("nodes_treeids", treeids); + test.AddAttribute("nodes_nodeids", nodeids); + test.AddAttribute("nodes_featureids", featureids); + test.AddAttribute("nodes_values", thresholds); + test.AddAttribute("nodes_modes", modes); + test.AddAttribute("target_treeids", target_treeids); + test.AddAttribute("target_nodeids", target_nodeids); + test.AddAttribute("target_ids", target_classids); + test.AddAttribute("target_weights", target_weights); + test.AddAttribute("base_values", base_values); + + test.AddAttribute("n_targets", (int64_t)2); + + if (aggFunction == "AVERAGE") { + test.AddAttribute("aggregate_function", "AVERAGE"); + } else if (aggFunction == "MIN") { + test.AddAttribute("aggregate_function", "MIN"); + } else if (aggFunction == "MAX") { + test.AddAttribute("aggregate_function", "MAX"); + } // default function is SUM + + //fill input data + test.AddInput("X", {8, 3}, X); + test.AddOutput("Y", {8, 2}, results); + test.Run(); +} + +TEST(MLOpTest, TreeRegressorMultiTargetAverage) { std::vector X = {1.f, 0.0f, 0.4f, 3.0f, 44.0f, -3.f, 12.0f, 12.9f, -312.f, 23.0f, 11.3f, -222.f, 23.0f, 11.3f, -222.f, 23.0f, 3311.3f, -222.f, 23.0f, 11.3f, -222.f, 43.0f, 413.3f, -114.f}; std::vector results = {1.33333333f, 29.f, 3.f, 14.f, 2.f, 23.f, 2.f, 23.f, 2.f, 23.f, 2.66666667f, 17.f, 2.f, 23.f, 3.f, 14.f}; + std::vector base_values{0.f, 0.f}; + GenTreeAndRunTest(X, base_values, results, "AVERAGE"); +} + +TEST(MLOpTest, TreeRegressorMultiTargetMin) { + std::vector X = {1.f, 0.0f, 0.4f, 3.0f, 44.0f, -3.f, 12.0f, 12.9f, -312.f, 23.0f, 11.3f, -222.f, 23.0f, 11.3f, -222.f, 23.0f, 3311.3f, -222.f, 23.0f, 11.3f, -222.f, 43.0f, 413.3f, -114.f}; + std::vector results = {5.f, 28.f, 8.f, 19.f, 7.f, 28.f, 7.f, 28.f, 7.f, 28.f, 7.f, 19.f, 7.f, 28.f, 8.f, 19.f}; + std::vector base_values{5.f, 5.f}; + GenTreeAndRunTest(X, base_values, results, "MIN"); +} + +TEST(MLOpTest, TreeRegressorMultiTargetMax) { + std::vector X = {1.f, 0.0f, 0.4f, 3.0f, 44.0f, -3.f, 12.0f, 12.9f, -312.f, 23.0f, 11.3f, -222.f, 23.0f, 11.3f, -222.f, 23.0f, 3311.3f, -222.f, 23.0f, 11.3f, -222.f, 43.0f, 413.3f, -114.f}; + std::vector results = {2.f, 41.f, 3.f, 14.f, 2.f, 23.f, 2.f, 23.f, 2.f, 23.f, 3.f, 23.f, 2.f, 23.f, 3.f, 14.f}; + std::vector base_values{0.f, 0.f}; + GenTreeAndRunTest(X, base_values, results, "MAX"); +} + +TEST(MLOpTest, TreeRegressorSingleTargetSum) { + OpTester test("TreeEnsembleRegressor", 1, onnxruntime::kMLDomain); + + //tree + std::vector lefts = {1, 0, 0, 1, 0, 0, 1, 0 ,0}; + std::vector rights = {2,0,0,2,0,0,2,0,0}; + std::vector treeids = {0,0,0,1,1,1,2,2,2}; + std::vector nodeids = {0,1,2,0,1,2,0,1,2}; + std::vector featureids = {0,0,0,0,0,0,1,0,0}; + std::vector thresholds = {1,0,0,0.5,0,0,0.5,0,0 }; + std::vector modes = {"BRANCH_LEQ", "LEAF", "LEAF", "BRANCH_LEQ", "LEAF", "LEAF", "BRANCH_LEQ", "LEAF", "LEAF"}; + + std::vector target_treeids = {0,0,1,1,2,2}; + std::vector target_nodeids = {1,2,1,2,1,2}; + std::vector target_classids = {0,0,0,0,0,0}; + std::vector target_weights = {33.33333f, 16.66666f, 33.33333f, -3.33333f, 16.66666f, -3.333333f}; + std::vector classes = {0, 1}; + + //test data + std::vector X = {0,1,1,1,2,0}; + std::vector results = {63.33333333f, 26.66666667f, 30.0f}; //add attributes test.AddAttribute("nodes_truenodeids", lefts); @@ -42,11 +112,12 @@ TEST(MLOpTest, TreeRegressorMultiTarget) { test.AddAttribute("target_ids", target_classids); test.AddAttribute("target_weights", target_weights); - test.AddAttribute("n_targets", (int64_t)2); - test.AddAttribute("aggregate_function", "AVERAGE"); + test.AddAttribute("n_targets", (int64_t)1); + // SUM aggregation by default -- no need to add explicitly + //fill input data - test.AddInput("X", {8, 3}, X); - test.AddOutput("Y", {8, 2}, results); + test.AddInput("X", {3, 2}, X); + test.AddOutput("Y", {3, 1}, results); test.Run(); }