saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Make operator TreeEnsemble 5x faster for batches of size 100.000 (#5965)

Browse source 
* improves processing time by 10
* extend coverage unit test coverage
* better implementation for the multi regression case
* better comment, keep parallelization by trees when not enough trees
This commit is contained in:
Xavier Dupré 2020-12-03 14:36:42 +01:00 committed by GitHub
parent 524b9fa899
commit 0acc3837ee
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 301 additions and 91 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
onnxruntime/core/providers/cpu/ml/tree_ensemble_classifier.cc
View file

@ -139,7 +139,7 @@ template <typename T>
TreeEnsembleClassifier<T>::TreeEnsembleClassifier(const OpKernelInfo& info)
: OpKernel(info),
tree_ensemble_(
100,
80,
50,
info.GetAttrOrDefault<std::string>("aggregate_function", "SUM"),
info.GetAttrsOrDefault<float>("base_values"),

224
onnxruntime/core/providers/cpu/ml/tree_ensemble_common.h
View file

@ -262,126 +262,180 @@ void TreeEnsembleCommon<ITYPE, OTYPE>::ComputeAgg(concurrency::ThreadPool* ttp,
const ITYPE* x_data = X->template Data<ITYPE>();
OTYPE* z_data = Z->template MutableData<OTYPE>();
int64_t* label_data = label == nullptr ? nullptr : label->template MutableData<int64_t>();
auto max_num_threads = concurrency::ThreadPool::DegreeOfParallelism(ttp);
if (n_targets_or_classes_ == 1) {
if (N == 1) {
ScoreValue<OTYPE> 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<ScoreValue<OTYPE>> scores_t(n_trees_, {0, 0});
} else { /* section B: 1 output, 1 row and enough trees to parallelize */
std::vector<ScoreValue<OTYPE>> scores(n_trees_, {0, 0});
concurrency::ThreadPool::TryBatchParallelFor(
ttp,
SafeInt<int32_t>(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<OTYPE> score;
size_t j;
} else if (N <= parallel_N_) { /* section C: 1 output, 2+ rows but not enough rows to parallelize */
ScoreValue<OTYPE> score;
size_t j;
for (int64_t i = 0; i < N; ++i) {
score = {0, 0};
for (j = 0; j < static_cast<size_t>(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<size_t>(n_trees_); ++j) {
agg.ProcessTreeNodePrediction1(score, *ProcessTreeNodeLeave(roots_[j], x_data + i * stride));
}
} else {
concurrency::ThreadPool::TryBatchParallelFor(
ttp,
SafeInt<int32_t>(N),
[this, &agg, x_data, z_data, stride, label_data](ptrdiff_t i) {
ScoreValue<OTYPE> score = {0, 0};
for (size_t j = 0; j < static_cast<size_t>(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<int32_t>(max_num_threads, SafeInt<int32_t>(n_trees_));
std::vector<ScoreValue<OTYPE>> 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<int32_t>(N),
[this, &agg, x_data, z_data, stride, label_data](ptrdiff_t i) {
ScoreValue<OTYPE> score = {0, 0};
for (size_t j = 0; j < static_cast<size_t>(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<ScoreValue<OTYPE>> 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<ScoreValue<OTYPE>> 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<int32_t>(concurrency::ThreadPool::DegreeOfParallelism(ttp), SafeInt<int32_t>(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<int32_t>(max_num_threads, SafeInt<int32_t>(n_trees_));
std::vector<std::vector<ScoreValue<OTYPE>>> scores(num_threads);
concurrency::ThreadPool::TrySimpleParallelFor(
ttp,
num_threads,
[this, &agg, &scores, &merge_mutex, num_threads, x_data](ptrdiff_t batch_num) {
std::vector<ScoreValue<OTYPE>> 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<OrtMutex> lock(merge_mutex);
agg.MergePrediction(scores, private_scores);
});
}
agg.FinalizeScores(scores, z_data, -1, label_data);
} else {
if (N <= parallel_N_) {
std::vector<ScoreValue<OTYPE>> scores(n_targets_or_classes_);
size_t j;
for (int64_t i = 0; i < N; ++i) {
std::fill(scores.begin(), scores.end(), ScoreValue<OTYPE>({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<int32_t>(concurrency::ThreadPool::DegreeOfParallelism(ttp), SafeInt<int32_t>(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<ScoreValue<OTYPE>> 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<OTYPE>({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<ScoreValue<OTYPE>> scores(n_targets_or_classes_);
size_t j;
for (int64_t i = 0; i < N; ++i) {
std::fill(scores.begin(), scores.end(), ScoreValue<OTYPE>({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<int32_t>(max_num_threads, SafeInt<int32_t>(n_trees_));
std::vector<std::vector<ScoreValue<OTYPE>>> 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<int32_t>(max_num_threads, SafeInt<int32_t>(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<ScoreValue<OTYPE>> 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<OTYPE>({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

2
onnxruntime/core/providers/cpu/ml/treeregressor.cc
View file

@ -24,7 +24,7 @@ template <typename T>
TreeEnsembleRegressor<T>::TreeEnsembleRegressor(const OpKernelInfo& info)
: OpKernel(info),
tree_ensemble_(
100,
80,
50,
info.GetAttrOrDefault<std::string>("aggregate_function", "SUM"),
info.GetAttrsOrDefault<float>("base_values"),

164
onnxruntime/test/providers/cpu/ml/treeregressor_test.cc
View file

@ -8,7 +8,31 @@ namespace onnxruntime {
namespace test {
template <typename T>
void GenTreeAndRunTest(const std::vector<T>& X, const std::vector<float>& base_values, const std::vector<float>& results, const std::string& aggFunction, bool one_obs = false) {
void _multiply_update_array(std::vector<T>& data, int n, T inc = 0) {
std::vector<T> 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<std::string>& data, int n) {
std::vector<std::string> 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 <typename T>
void GenTreeAndRunTest(const std::vector<T>& X, const std::vector<float>& base_values, const std::vector<float>& 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<T>& X, const std::vector<float>& base_v
std::vector<float> 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<int64_t> 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<T>& X, const std::vector<float>& base_v
} // default function is SUM
//fill input data
std::vector<T> xn;
std::vector<float> yn;
if (one_obs) {
auto X1 = X;
auto results1 = results;
@ -58,13 +99,69 @@ void GenTreeAndRunTest(const std::vector<T>& X, const std::vector<float>& base_v
results1.resize(2);
test.AddInput<T>("X", {1, 3}, X1);
test.AddOutput<float>("Y", {1, 2}, results1);
} else {
} else if (n_obs == 8) {
test.AddInput<T>("X", {8, 3}, X);
test.AddOutput<float>("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<T>("X", {n_obs, 3}, xn);
test.AddOutput<float>("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<float> 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<float> 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<float> base_values{0.f, 0.f};
GenTreeAndRunTest(X, base_values, results, "AVERAGE", true, 8, 1); // section A2
}
TEST(MLOpTest, TreeRegressorMultiTargetBatchTreeB2) {
std::vector<float> 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<float> 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<float> base_values{0.f, 0.f};
GenTreeAndRunTest(X, base_values, results, "AVERAGE", true, 8, 130); // section B2
}
TEST(MLOpTest, TreeRegressorMultiTargetBatchTreeC2) {
std::vector<float> 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<float> 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<float> base_values{0.f, 0.f};
GenTreeAndRunTest(X, base_values, results, "AVERAGE", false, 200, 130); // section C2
}
TEST(MLOpTest, TreeRegressorMultiTargetBatchTreeD2) {
std::vector<float> 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<float> 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<float> 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<float> 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<float> 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<float> base_values{0.f, 0.f};
GenTreeAndRunTest(X, base_values, results, "AVERAGE", false, 200, 1); // section E2
}
TEST(MLOpTest, TreeRegressorMultiTargetAverage) {
std::vector<float> 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<float> 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<double>(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<float> target_weights = {33.33333f, 16.66666f, 33.33333f, -3.33333f, 16.66666f, -3.333333f};
std::vector<int64_t> 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<float> 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<float> xn, yn;
if (one_obs) {
ASSERT_TRUE(n_obs == 3);
auto X1 = X;
auto results1 = results;
X1.resize(2);
results1.resize(1);
test.AddInput<float>("X", {1, 2}, X1);
test.AddOutput<float>("Y", {1, 1}, results1);
} else {
} else if (n_obs == 3) {
test.AddInput<float>("X", {3, 2}, X);
test.AddOutput<float>("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<float>("X", {n_obs, 2}, xn);
test.AddOutput<float>("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);