diff --git a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h index a88953832f..f951d7ccbb 100644 --- a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h +++ b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h @@ -133,8 +133,8 @@ TreeEnsembleCommon::TreeEnsembleCommon(int parallel_tree, int para node.hitrates = i < nodes_hitrates.size() ? nodes_hitrates[i] : -1; node.mode = cmodes[i]; node.is_not_leaf = node.mode != NODE_MODE::LEAF; - node.truenode = NULL; // nodes_truenodeids[i]; - node.falsenode = NULL; // nodes_falsenodeids[i]; + node.truenode = nullptr; // nodes_truenodeids[i]; + node.falsenode = nullptr; // nodes_falsenodeids[i]; node.missing_tracks = i < static_cast(nodes_missing_value_tracks_true.size()) ? (nodes_missing_value_tracks_true[i] == 1 ? MissingTrack::TRUE @@ -166,7 +166,7 @@ TreeEnsembleCommon::TreeEnsembleCommon(int parallel_tree, int para ORT_THROW("One falsenode is pointing either to itself, either to another tree."); } } else - it->truenode = NULL; + it->truenode = nullptr; coor.node_id = static_cast(nodes_falsenodeids[i]); found = idi.find(coor); @@ -180,7 +180,7 @@ TreeEnsembleCommon::TreeEnsembleCommon(int parallel_tree, int para ORT_THROW("One falsenode is pointing either to itself, either to another tree."); } } else - it->falsenode = NULL; + it->falsenode = nullptr; } int64_t previous = -1; @@ -215,7 +215,8 @@ TreeEnsembleCommon::TreeEnsembleCommon(int parallel_tree, int para } template -void TreeEnsembleCommon::compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, Tensor* label) const { +void TreeEnsembleCommon::compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, + Tensor* label) const { switch (aggregate_function_) { case AGGREGATE_FUNCTION::AVERAGE: ComputeAgg( @@ -252,31 +253,35 @@ void TreeEnsembleCommon::compute(concurrency::ThreadPool* ttp, con template template -void TreeEnsembleCommon::ComputeAgg(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, Tensor* label, const AGG& agg) const { +void TreeEnsembleCommon::ComputeAgg(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, + Tensor* label, const AGG& agg) const { int64_t stride = X->Shape().NumDimensions() == 1 ? X->Shape()[0] : X->Shape()[1]; int64_t N = X->Shape().NumDimensions() == 1 ? 1 : X->Shape()[0]; const ITYPE* x_data = X->template Data(); OTYPE* z_data = Z->template MutableData(); - int64_t* label_data = label == NULL ? NULL : label->template MutableData(); + int64_t* label_data = label == nullptr ? nullptr : label->template MutableData(); if (n_targets_or_classes_ == 1) { if (N == 1) { ScoreValue score = {0, 0}; if (n_trees_ <= parallel_tree_) { - for (int64_t j = 0; j < n_trees_; ++j) + for (int64_t j = 0; j < n_trees_; ++j) { agg.ProcessTreeNodePrediction1(score, *ProcessTreeNodeLeave(roots_[j], x_data)); + } } else { std::vector> scores_t(n_trees_, {0, 0}); concurrency::ThreadPool::TryBatchParallelFor( ttp, - static_cast(this->n_trees_), - [&](ptrdiff_t j) { - agg.ProcessTreeNodePrediction1(scores_t[j], *ProcessTreeNodeLeave(this->roots_[j], x_data)); + SafeInt(n_trees_), + [this, &scores_t, &agg, x_data](ptrdiff_t j) { + agg.ProcessTreeNodePrediction1(scores_t[j], *ProcessTreeNodeLeave(roots_[j], x_data)); }, 0); - for (auto it = scores_t.cbegin(); it != scores_t.cend(); ++it) + + for (auto it = scores_t.cbegin(); it != scores_t.cend(); ++it) { agg.MergePrediction1(score, *it); + } } agg.FinalizeScores1(z_data, score, label_data); @@ -287,21 +292,25 @@ void TreeEnsembleCommon::ComputeAgg(concurrency::ThreadPool* ttp, for (int64_t i = 0; i < N; ++i) { score = {0, 0}; - for (j = 0; j < static_cast(n_trees_); ++j) + for (j = 0; j < static_cast(n_trees_); ++j) { agg.ProcessTreeNodePrediction1(score, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + } + agg.FinalizeScores1(z_data + i * n_targets_or_classes_, score, - label_data == NULL ? NULL : (label_data + i)); + label_data == nullptr ? nullptr : (label_data + i)); } } else { concurrency::ThreadPool::TryBatchParallelFor( ttp, - static_cast(N), - [&](ptrdiff_t i) { + SafeInt(N), + [this, &agg, x_data, z_data, stride, label_data](ptrdiff_t i) { ScoreValue score = {0, 0}; - for (size_t j = 0; j < static_cast(n_trees_); ++j) + for (size_t j = 0; j < static_cast(n_trees_); ++j) { agg.ProcessTreeNodePrediction1(score, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + } + agg.FinalizeScores1(z_data + i * n_targets_or_classes_, score, - label_data == NULL ? NULL : (label_data + i)); + label_data == nullptr ? nullptr : (label_data + i)); }, 0); } @@ -309,30 +318,32 @@ void TreeEnsembleCommon::ComputeAgg(concurrency::ThreadPool* ttp, } else { if (N == 1) { std::vector> scores(n_targets_or_classes_, {0, 0}); - if (n_trees_ <= parallel_tree_) { - for (int64_t j = 0; j < n_trees_; ++j) + for (int64_t j = 0; j < n_trees_; ++j) { agg.ProcessTreeNodePrediction(scores, *ProcessTreeNodeLeave(roots_[j], x_data)); - agg.FinalizeScores(scores, z_data, -1, label_data); + } } else { - auto nth = n_trees_ * 2 / concurrency::ThreadPool::NumThreads(ttp); - if (n_trees_ % nth != 0) - ++nth; - concurrency::ThreadPool::TryBatchParallelFor( + // split the work into one block per thread so we can re-use the 'private_scores' vector as much as possible + // TODO: Refine the number of threads used + auto num_threads = std::min(concurrency::ThreadPool::NumThreads(ttp), SafeInt(n_trees_)); + std::mutex merge_mutex; + concurrency::ThreadPool::TrySimpleParallelFor( ttp, - static_cast(nth), - [&](ptrdiff_t th) { + num_threads, + [this, &agg, &scores, &merge_mutex, num_threads, x_data](ptrdiff_t batch_num) { std::vector> private_scores(n_targets_or_classes_, {0, 0}); - auto end = nth * (th + 1); - end = end < N ? end : N; - for (int64_t j = nth * th; j < end; ++j) { + int64_t j = batch_num * n_trees_ / num_threads; + int64_t end = (batch_num + 1) * n_trees_ / num_threads; + for (; j < end; ++j) { agg.ProcessTreeNodePrediction(private_scores, *ProcessTreeNodeLeave(roots_[j], x_data)); } + + std::lock_guard lock(merge_mutex); agg.MergePrediction(scores, private_scores); - }, - 0); - agg.FinalizeScores(scores, z_data, -1, label_data); + }); } + + agg.FinalizeScores(scores, z_data, -1, label_data); } else { if (N <= parallel_N_) { std::vector> scores(n_targets_or_classes_); @@ -340,35 +351,36 @@ void TreeEnsembleCommon::ComputeAgg(concurrency::ThreadPool* ttp, for (int64_t i = 0; i < N; ++i) { std::fill(scores.begin(), scores.end(), ScoreValue({0, 0})); - for (j = 0; j < roots_.size(); ++j) + for (j = 0; j < roots_.size(); ++j) { agg.ProcessTreeNodePrediction(scores, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); - agg.FinalizeScores(scores, - z_data + i * n_targets_or_classes_, -1, - label_data == NULL ? NULL : (label_data + i)); - ORT_ENFORCE((int64_t)scores.size() == n_targets_or_classes_); + } + + agg.FinalizeScores(scores, z_data + i * n_targets_or_classes_, -1, + label_data == nullptr ? nullptr : (label_data + i)); } } else { - auto nth = N * 2 / concurrency::ThreadPool::NumThreads(ttp); - if (N % nth != 0) - ++nth; - concurrency::ThreadPool::TryBatchParallelFor( + // split the work into one block per thread so we can re-use the 'scores' vector as much as possible + // TODO: Refine the number of threads used. + auto num_threads = std::min(concurrency::ThreadPool::NumThreads(ttp), SafeInt(N)); + concurrency::ThreadPool::TrySimpleParallelFor( ttp, - static_cast(nth), - [&](ptrdiff_t th) { + num_threads, + [this, &agg, num_threads, x_data, z_data, label_data, N, stride](ptrdiff_t batch_num) { size_t j; std::vector> scores(n_targets_or_classes_); - auto end = nth * (th + 1); - end = end < N ? end : N; - for (int64_t i = nth * th; i < end; ++i) { + int64_t i = batch_num * N / num_threads; + int64_t end = (batch_num + 1) * N / num_threads; + for (; i < end; ++i) { std::fill(scores.begin(), scores.end(), ScoreValue({0, 0})); - for (j = 0; j < roots_.size(); ++j) + for (j = 0; j < roots_.size(); ++j) { agg.ProcessTreeNodePrediction(scores, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + } + agg.FinalizeScores(scores, z_data + i * n_targets_or_classes_, -1, - label_data == NULL ? NULL : (label_data + i)); + label_data == nullptr ? nullptr : (label_data + i)); } - }, - 0); + }); } } } @@ -537,25 +549,27 @@ TreeEnsembleCommonClassifier::TreeEnsembleCommonClassifier( const std::vector& class_treeids, const std::vector& class_weights, const std::vector& classlabels_strings, - const std::vector& classlabels_int64s) : TreeEnsembleCommon(parallel_tree, - parallel_N, - aggregate_function, - base_values, - classlabels_strings.size() == 0 ? classlabels_int64s.size() : classlabels_strings.size(), - nodes_falsenodeids, - nodes_featureids, - nodes_hitrates, - nodes_missing_value_tracks_true, - nodes_modes, - nodes_nodeids, - nodes_treeids, - nodes_truenodeids, - nodes_values, - post_transform, - class_ids, - class_nodeids, - class_treeids, - class_weights) { + const std::vector& classlabels_int64s) + : TreeEnsembleCommon(parallel_tree, + parallel_N, + aggregate_function, + base_values, + classlabels_strings.size() == 0 ? classlabels_int64s.size() + : classlabels_strings.size(), + nodes_falsenodeids, + nodes_featureids, + nodes_hitrates, + nodes_missing_value_tracks_true, + nodes_modes, + nodes_nodeids, + nodes_treeids, + nodes_truenodeids, + nodes_values, + post_transform, + class_ids, + class_nodeids, + class_treeids, + class_weights) { classlabels_strings_ = classlabels_strings; classlabels_int64s_ = classlabels_int64s; @@ -575,7 +589,8 @@ TreeEnsembleCommonClassifier::TreeEnsembleCommonClassifier( } template -void TreeEnsembleCommonClassifier::compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, Tensor* label) const { +void TreeEnsembleCommonClassifier::compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, + Tensor* label) const { if (classlabels_strings_.size() == 0) { this->ComputeAgg( ttp, X, Z, label, diff --git a/onnxruntime/test/providers/cpu/ml/tree_ensembler_classifier_test.cc b/onnxruntime/test/providers/cpu/ml/tree_ensembler_classifier_test.cc index fdad2ed54c..f0ec922919 100644 --- a/onnxruntime/test/providers/cpu/ml/tree_ensembler_classifier_test.cc +++ b/onnxruntime/test/providers/cpu/ml/tree_ensembler_classifier_test.cc @@ -54,6 +54,51 @@ TEST(MLOpTest, TreeEnsembleClassifier) { test.Run(); } +TEST(MLOpTest, TreeEnsembleClassifier_N1) { + OpTester test("TreeEnsembleClassifier", 1, onnxruntime::kMLDomain); + + std::vector lefts = {1, -1, 3, -1, -1, 1, -1, 3, 4, -1, -1, -1, 1, 2, -1, 4, -1, -1, -1}; + std::vector rights = {2, -1, 4, -1, -1, 2, -1, 6, 5, -1, -1, -1, 6, 3, -1, 5, -1, -1, -1}; + std::vector treeids = {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2}; + std::vector nodeids = {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6}; + std::vector featureids = {2, -2, 0, -2, -2, 0, -2, 2, 1, -2, -2, -2, 0, 2, -2, 1, -2, -2, -2}; + std::vector thresholds = {-172.f, -2.f, 2.5f, -2.f, -2.f, 1.5f, -2.f, -62.5f, 213.09999084f, + -2.f, -2.f, -2.f, 27.5f, -172.f, -2.f, 8.10000038f, -2.f, -2.f, -2.f}; + std::vector modes = {"BRANCH_LEQ", "LEAF", "BRANCH_LEQ", "LEAF", "LEAF", "BRANCH_LEQ", + "LEAF", "BRANCH_LEQ", "BRANCH_LEQ", "LEAF", "LEAF", "LEAF", + "BRANCH_LEQ", "BRANCH_LEQ", "LEAF", "BRANCH_LEQ", "LEAF", "LEAF", "LEAF"}; + std::vector class_treeids = {0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2}; + std::vector class_nodeids = {1, 3, 4, 1, 4, 5, 6, 2, 4, 5, 6}; + std::vector class_classids = {2, 0, 1, 0, 2, 3, 1, 2, 0, 1, 3}; + std::vector class_weights = {1.f, 4.f, 1.f, 2.f, 1.f, 1.f, 2.f, 1.f, 1.f, 1.f, 3.f}; + std::vector classes = {0, 1, 2, 3}; + std::vector X = {1.f, 0.0f, 0.4f}; + std::vector results = {0}; + std::vector scores{7, 0, 0, 0}; + std::vector probs = {}; + std::vector log_probs = {}; + + //define the context of the operator call + const int N = 1; + 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("class_treeids", class_treeids); + test.AddAttribute("class_nodeids", class_nodeids); + test.AddAttribute("class_ids", class_classids); + test.AddAttribute("class_weights", class_weights); + test.AddAttribute("classlabels_int64s", classes); + + test.AddInput("X", {N, 3}, X); + test.AddOutput("Y", {N}, results); + test.AddOutput("Z", {N, static_cast(classes.size())}, scores); + test.Run(); +} + TEST(MLOpTest, TreeEnsembleClassifierLabels) { OpTester test("TreeEnsembleClassifier", 1, onnxruntime::kMLDomain);