diff --git a/onnxruntime/core/providers/cpu/ml/tree_ensemble_classifier.cc b/onnxruntime/core/providers/cpu/ml/tree_ensemble_classifier.cc index f4512e5df0..216bacd607 100644 --- a/onnxruntime/core/providers/cpu/ml/tree_ensemble_classifier.cc +++ b/onnxruntime/core/providers/cpu/ml/tree_ensemble_classifier.cc @@ -139,7 +139,7 @@ template TreeEnsembleClassifier::TreeEnsembleClassifier(const OpKernelInfo& info) : OpKernel(info), tree_ensemble_( - 100, + 80, 50, info.GetAttrOrDefault("aggregate_function", "SUM"), info.GetAttrsOrDefault("base_values"), diff --git a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h index 0dbdf544bc..4a06f5a7a2 100644 --- a/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h +++ b/onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h @@ -262,126 +262,180 @@ void TreeEnsembleCommon::ComputeAgg(concurrency::ThreadPool* ttp, const ITYPE* x_data = X->template Data(); OTYPE* z_data = Z->template MutableData(); int64_t* label_data = label == nullptr ? nullptr : label->template MutableData(); + auto max_num_threads = concurrency::ThreadPool::DegreeOfParallelism(ttp); if (n_targets_or_classes_ == 1) { if (N == 1) { ScoreValue score = {0, 0}; - if (n_trees_ <= parallel_tree_) { + if (n_trees_ <= parallel_tree_) { /* section A: 1 output, 1 row and not enough trees to parallelize */ 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}); + } else { /* section B: 1 output, 1 row and enough trees to parallelize */ + std::vector> scores(n_trees_, {0, 0}); concurrency::ThreadPool::TryBatchParallelFor( ttp, SafeInt(n_trees_), - [this, &scores_t, &agg, x_data](ptrdiff_t j) { - agg.ProcessTreeNodePrediction1(scores_t[j], *ProcessTreeNodeLeave(roots_[j], x_data)); + [this, &scores, &agg, x_data](ptrdiff_t j) { + agg.ProcessTreeNodePrediction1(scores[j], *ProcessTreeNodeLeave(roots_[j], x_data)); }, 0); - for (auto it = scores_t.cbegin(); it != scores_t.cend(); ++it) { + for (auto it = scores.cbegin(); it != scores.cend(); ++it) { agg.MergePrediction1(score, *it); } } - agg.FinalizeScores1(z_data, score, label_data); - } else { - if (N <= parallel_N_) { - ScoreValue score; - size_t j; + } else if (N <= parallel_N_) { /* section C: 1 output, 2+ rows but not enough rows to parallelize */ + ScoreValue score; + size_t j; - for (int64_t i = 0; i < N; ++i) { - score = {0, 0}; - 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 == nullptr ? nullptr : (label_data + i)); + for (int64_t i = 0; i < N; ++i) { + score = {0, 0}; + for (j = 0; j < static_cast(n_trees_); ++j) { + agg.ProcessTreeNodePrediction1(score, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); } - } else { - concurrency::ThreadPool::TryBatchParallelFor( - ttp, - 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) { - agg.ProcessTreeNodePrediction1(score, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); - } - agg.FinalizeScores1(z_data + i * n_targets_or_classes_, score, - label_data == nullptr ? nullptr : (label_data + i)); - }, - 0); + agg.FinalizeScores1(z_data + i, score, + label_data == nullptr ? nullptr : (label_data + i)); } + } else if (n_trees_ > max_num_threads) { /* section D: 1 output, 2+ rows and enough trees to parallelize */ + auto num_threads = std::min(max_num_threads, SafeInt(n_trees_)); + std::vector> scores(num_threads * N); + concurrency::ThreadPool::TrySimpleParallelFor( + ttp, + num_threads, + [this, &agg, &scores, num_threads, x_data, N, stride](ptrdiff_t batch_num) { + auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, this->n_trees_); + for (int64_t i = 0; i < N; ++i) { + scores[batch_num * N + i] = {0, 0}; + } + for (auto j = work.start; j < work.end; ++j) { + for (int64_t i = 0; i < N; ++i) { + agg.ProcessTreeNodePrediction1(scores[batch_num * N + i], *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + } + } + }); + + concurrency::ThreadPool::TrySimpleParallelFor( + ttp, + num_threads, + [&agg, &scores, num_threads, label_data, z_data, N](ptrdiff_t batch_num) { + auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, N); + for (auto i = work.start; i < work.end; ++i) { + for (int64_t j = 1; j < num_threads; ++j) { + agg.MergePrediction1(scores[i], scores[j * N + i]); + } + agg.FinalizeScores1(z_data + i, scores[i], + label_data == nullptr ? nullptr : (label_data + i)); + } + }); + } else { /* section E: 1 output, 2+ rows, parallelization by rows */ + concurrency::ThreadPool::TryBatchParallelFor( + ttp, + 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) { + agg.ProcessTreeNodePrediction1(score, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + } + + agg.FinalizeScores1(z_data + i, score, + label_data == nullptr ? nullptr : (label_data + i)); + }, + 0); } } else { - if (N == 1) { - std::vector> scores(n_targets_or_classes_, {0, 0}); - if (n_trees_ <= parallel_tree_) { + if (N == 1) { /* section A2: 2+ outputs, 1 row, not enough trees to parallelize */ + if (n_trees_ <= parallel_tree_) { /* section A2 */ + std::vector> scores(n_targets_or_classes_, {0, 0}); for (int64_t j = 0; j < n_trees_; ++j) { agg.ProcessTreeNodePrediction(scores, *ProcessTreeNodeLeave(roots_[j], x_data)); } - } else { - // 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::DegreeOfParallelism(ttp), SafeInt(n_trees_)); - OrtMutex merge_mutex; + agg.FinalizeScores(scores, z_data, -1, label_data); + } else { /* section B2: 2+ outputs, 1 row, enough trees to parallelize */ + auto num_threads = std::min(max_num_threads, SafeInt(n_trees_)); + std::vector>> scores(num_threads); concurrency::ThreadPool::TrySimpleParallelFor( ttp, 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}); + [this, &agg, &scores, num_threads, x_data](ptrdiff_t batch_num) { + scores[batch_num].resize(n_targets_or_classes_, {0, 0}); auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, n_trees_); for (auto j = work.start; j < work.end; ++j) { - agg.ProcessTreeNodePrediction(private_scores, *ProcessTreeNodeLeave(roots_[j], x_data)); + agg.ProcessTreeNodePrediction(scores[batch_num], *ProcessTreeNodeLeave(roots_[j], x_data)); } - - std::lock_guard lock(merge_mutex); - agg.MergePrediction(scores, private_scores); }); - } - - agg.FinalizeScores(scores, z_data, -1, label_data); - } else { - if (N <= parallel_N_) { - std::vector> scores(n_targets_or_classes_); - size_t j; - - for (int64_t i = 0; i < N; ++i) { - std::fill(scores.begin(), scores.end(), ScoreValue({0, 0})); - 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 == nullptr ? nullptr : (label_data + i)); + for (size_t i = 1; i < scores.size(); ++i) { + agg.MergePrediction(scores[0], scores[i]); } - } else { - // 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::DegreeOfParallelism(ttp), SafeInt(N)); - concurrency::ThreadPool::TrySimpleParallelFor( - ttp, - 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 work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, N); - - for (auto i = work.start; i < work.end; ++i) { - std::fill(scores.begin(), scores.end(), ScoreValue({0, 0})); - 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 == nullptr ? nullptr : (label_data + i)); - } - }); + agg.FinalizeScores(scores[0], z_data, -1, label_data); } + } else if (N <= parallel_N_) { /* section C2: 2+ outputs, 2+ rows, not enough rows to parallelize */ + std::vector> scores(n_targets_or_classes_); + size_t j; + + for (int64_t i = 0; i < N; ++i) { + std::fill(scores.begin(), scores.end(), ScoreValue({0, 0})); + 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 == nullptr ? nullptr : (label_data + i)); + } + } else if (n_trees_ >= max_num_threads) { /* section: D2: 2+ outputs, 2+ rows, enough trees to parallelize*/ + auto num_threads = std::min(max_num_threads, SafeInt(n_trees_)); + std::vector>> scores(num_threads * N); + concurrency::ThreadPool::TrySimpleParallelFor( + ttp, + num_threads, + [this, &agg, &scores, num_threads, x_data, N, stride](ptrdiff_t batch_num) { + auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, this->n_trees_); + for (int64_t i = 0; i < N; ++i) { + scores[batch_num * N + i].resize(n_targets_or_classes_, {0, 0}); + } + for (auto j = work.start; j < work.end; ++j) { + for (int64_t i = 0; i < N; ++i) { + agg.ProcessTreeNodePrediction(scores[batch_num * N + i], *ProcessTreeNodeLeave(roots_[j], x_data + i * stride)); + } + } + }); + + concurrency::ThreadPool::TrySimpleParallelFor( + ttp, + num_threads, + [this, &agg, &scores, num_threads, label_data, z_data, N](ptrdiff_t batch_num) { + auto work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, N); + for (auto i = work.start; i < work.end; ++i) { + for (int64_t j = 1; j < num_threads; ++j) { + agg.MergePrediction(scores[i], scores[j * N + i]); + } + agg.FinalizeScores(scores[i], z_data + i * this->n_targets_or_classes_, -1, + label_data == nullptr ? nullptr : (label_data + i)); + } + }); + } else { /* section E2: 2+ outputs, 2+ rows, parallelization by rows */ + auto num_threads = std::min(max_num_threads, SafeInt(N)); + concurrency::ThreadPool::TrySimpleParallelFor( + ttp, + 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 work = concurrency::ThreadPool::PartitionWork(batch_num, num_threads, N); + + for (auto i = work.start; i < work.end; ++i) { + std::fill(scores.begin(), scores.end(), ScoreValue({0, 0})); + 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 == nullptr ? nullptr : (label_data + i)); + } + }); } } } // namespace detail diff --git a/onnxruntime/core/providers/cpu/ml/treeregressor.cc b/onnxruntime/core/providers/cpu/ml/treeregressor.cc index cfee2ccae8..960e4fcf97 100644 --- a/onnxruntime/core/providers/cpu/ml/treeregressor.cc +++ b/onnxruntime/core/providers/cpu/ml/treeregressor.cc @@ -24,7 +24,7 @@ template TreeEnsembleRegressor::TreeEnsembleRegressor(const OpKernelInfo& info) : OpKernel(info), tree_ensemble_( - 100, + 80, 50, info.GetAttrOrDefault("aggregate_function", "SUM"), info.GetAttrsOrDefault("base_values"), diff --git a/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc b/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc index 87a3ecde9d..e6effd1427 100644 --- a/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc +++ b/onnxruntime/test/providers/cpu/ml/treeregressor_test.cc @@ -8,7 +8,31 @@ namespace onnxruntime { namespace test { template -void GenTreeAndRunTest(const std::vector& X, const std::vector& base_values, const std::vector& results, const std::string& aggFunction, bool one_obs = false) { +void _multiply_update_array(std::vector& data, int n, T inc = 0) { + std::vector copy = data; + data.resize(copy.size() * n); + T cst = 0; + for (int i = 0; i < n; ++i) { + for (size_t j = 0; j < copy.size(); ++j) { + data[j + i * copy.size()] = copy[j] + cst; + } + cst += inc; + } +} + +void _multiply_update_array_string(std::vector& data, int n) { + std::vector copy = data; + data.resize(copy.size() * n); + for (int i = 0; i < n; ++i) { + for (size_t j = 0; j < copy.size(); ++j) { + data[j + i * copy.size()] = copy[j]; + } + } +} + +template +void GenTreeAndRunTest(const std::vector& X, const std::vector& base_values, const std::vector& results, const std::string& aggFunction, + bool one_obs = false, int64_t n_obs = 8, int n_trees = 1) { OpTester test("TreeEnsembleRegressor", 1, onnxruntime::kMLDomain); //tree @@ -26,6 +50,21 @@ void GenTreeAndRunTest(const std::vector& X, const std::vector& base_v 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}; + if (n_trees > 1) { + // Multiplies the number of trees to test the parallelization by trees. + _multiply_update_array(lefts, n_trees); + _multiply_update_array(rights, n_trees); + _multiply_update_array(treeids, n_trees, (int64_t)3); + _multiply_update_array(nodeids, n_trees); + _multiply_update_array(featureids, n_trees); + _multiply_update_array(thresholds, n_trees); + _multiply_update_array_string(modes, n_trees); + _multiply_update_array(target_treeids, n_trees, (int64_t)3); + _multiply_update_array(target_nodeids, n_trees); + _multiply_update_array(target_classids, n_trees); + _multiply_update_array(target_weights, n_trees); + } + //add attributes test.AddAttribute("nodes_truenodeids", lefts); test.AddAttribute("nodes_falsenodeids", rights); @@ -51,6 +90,8 @@ void GenTreeAndRunTest(const std::vector& X, const std::vector& base_v } // default function is SUM //fill input data + std::vector xn; + std::vector yn; if (one_obs) { auto X1 = X; auto results1 = results; @@ -58,13 +99,69 @@ void GenTreeAndRunTest(const std::vector& X, const std::vector& base_v results1.resize(2); test.AddInput("X", {1, 3}, X1); test.AddOutput("Y", {1, 2}, results1); - } else { + } else if (n_obs == 8) { test.AddInput("X", {8, 3}, X); test.AddOutput("Y", {8, 2}, results); + } else { + int64_t i; + size_t k; + ASSERT_TRUE(n_obs % 8 == 0); + xn.resize(n_obs * 3); + yn.resize(n_obs * 2); + for (i = 0; i < n_obs; i += 8) { + for (k = 0; k < 24; ++k) { + xn[i * 3 + k] = X[k]; + } + for (k = 0; k < 16; ++k) { + yn[i * 2 + k] = results[k]; + } + } + ASSERT_TRUE(i == n_obs); + test.AddInput("X", {n_obs, 3}, xn); + test.AddOutput("Y", {n_obs, 2}, yn); } + test.Run(); } // namespace test +TEST(MLOpTest, TreeRegressorMultiTargetBatchTreeA2) { + // TreeEnsemble implements different paths depending on n_trees or N. + // This test and the next ones go through all sections for multi-targets. + 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", true, 8, 1); // section A2 +} + +TEST(MLOpTest, TreeRegressorMultiTargetBatchTreeB2) { + 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", true, 8, 130); // section B2 +} + +TEST(MLOpTest, TreeRegressorMultiTargetBatchTreeC2) { + 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", false, 200, 130); // section C2 +} + +TEST(MLOpTest, TreeRegressorMultiTargetBatchTreeD2) { + 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", true, 8, 30); // section D2 + GenTreeAndRunTest(X, base_values, results, "AVERAGE", false, 200, 30); // section D2 +} + +TEST(MLOpTest, TreeRegressorMultiTargetBatchTreeE2) { + 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", false, 200, 1); // section E2 +} + 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}; @@ -97,7 +194,7 @@ TEST(MLOpTest, TreeRegressorMultiTargetMaxDouble) { GenTreeAndRunTest(X, base_values, results, "MAX", true); } -void GenTreeAndRunTest1(const std::string& aggFunction, bool one_obs) { +void GenTreeAndRunTest1(const std::string& aggFunction, bool one_obs, int64_t n_obs = 3, int n_trees = 1) { OpTester test("TreeEnsembleRegressor", 1, onnxruntime::kMLDomain); //tree @@ -115,6 +212,21 @@ void GenTreeAndRunTest1(const std::string& aggFunction, bool one_obs) { std::vector target_weights = {33.33333f, 16.66666f, 33.33333f, -3.33333f, 16.66666f, -3.333333f}; std::vector classes = {0, 1}; + if (n_trees > 1) { + // Multiplies the number of trees to test the parallelization by trees. + _multiply_update_array(lefts, n_trees); + _multiply_update_array(rights, n_trees); + _multiply_update_array(treeids, n_trees, (int64_t)3); + _multiply_update_array(nodeids, n_trees); + _multiply_update_array(featureids, n_trees); + _multiply_update_array(thresholds, n_trees); + _multiply_update_array_string(modes, n_trees); + _multiply_update_array(target_treeids, n_trees, (int64_t)3); + _multiply_update_array(target_nodeids, n_trees); + _multiply_update_array(target_classids, n_trees); + _multiply_update_array(target_weights, n_trees); + } + std::vector results; if (aggFunction == "AVERAGE") { test.AddAttribute("aggregate_function", "AVERAGE"); @@ -149,16 +261,32 @@ void GenTreeAndRunTest1(const std::string& aggFunction, bool one_obs) { // SUM aggregation by default -- no need to add explicitly //fill input data + std::vector xn, yn; if (one_obs) { + ASSERT_TRUE(n_obs == 3); auto X1 = X; auto results1 = results; X1.resize(2); results1.resize(1); test.AddInput("X", {1, 2}, X1); test.AddOutput("Y", {1, 1}, results1); - } else { + } else if (n_obs == 3) { test.AddInput("X", {3, 2}, X); test.AddOutput("Y", {3, 1}, results); + } else { + ASSERT_TRUE(n_obs % 3 == 0); + xn.resize(n_obs * 2); + yn.resize(n_obs); + for (int64_t i = 0; i < n_obs; i += 3) { + for (size_t k = 0; k < 6; ++k) { + xn[i * 2 + k] = X[k]; + } + for (size_t k = 0; k < 3; ++k) { + yn[i + k] = results[k]; + } + } + test.AddInput("X", {n_obs, 2}, xn); + test.AddOutput("Y", {n_obs, 1}, yn); } test.Run(); } @@ -168,6 +296,34 @@ TEST(MLOpTest, TreeRegressorSingleTargetSum) { GenTreeAndRunTest1("SUM", true); } +TEST(MLOpTest, TreeRegressorSingleTargetSumBatch) { + GenTreeAndRunTest1("SUM", false, 201); + GenTreeAndRunTest1("SUM", false, 40002); +} + +TEST(MLOpTest, TreeRegressorSingleTargetBatchTreeA) { + // TreeEnsemble implements different paths depending on n_trees or N. + // This test and the next ones goe through all sections for one target. + GenTreeAndRunTest1("SUM", true, 3, 1); // section A +} + +TEST(MLOpTest, TreeRegressorSingleTargetBatchTreeB) { + GenTreeAndRunTest1("AVERAGE", true, 3, 30); // section B +} + +TEST(MLOpTest, TreeRegressorSingleTargetBatchTreeC) { + GenTreeAndRunTest1("AVERAGE", false, 3, 1); // section C +} + +TEST(MLOpTest, TreeRegressorSingleTargetBatchTreeD) { + GenTreeAndRunTest1("AVERAGE", false, 201, 30); // section D + GenTreeAndRunTest1("AVERAGE", false, 201, 130); // section D +} + +TEST(MLOpTest, TreeRegressorSingleTargetBatchTreeE) { + GenTreeAndRunTest1("AVERAGE", false, 201, 1); // section E +} + TEST(MLOpTest, TreeRegressorSingleTargetAverage) { GenTreeAndRunTest1("AVERAGE", false); GenTreeAndRunTest1("AVERAGE", true);