mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-26 19:52:38 +00:00
Fix tree ensemble threading bug (#3778)
* Fix first instance where the calculation for TryBatchParallelFor was incorrect. Other usages need to be validated. * Fix some other usages of the threadpool.
This commit is contained in:
Scott McKay
GitHub
parent
2fc3984e70
commit
11b819054b
2 changed files with 131 additions and 71 deletions
|
|
@ -133,8 +133,8 @@ TreeEnsembleCommon<ITYPE, OTYPE>::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<size_t>(nodes_missing_value_tracks_true.size())
|
||||
? (nodes_missing_value_tracks_true[i] == 1
|
||||
? MissingTrack::TRUE
|
||||
|
|
@ -166,7 +166,7 @@ TreeEnsembleCommon<ITYPE, OTYPE>::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<int>(nodes_falsenodeids[i]);
|
||||
found = idi.find(coor);
|
||||
|
|
@ -180,7 +180,7 @@ TreeEnsembleCommon<ITYPE, OTYPE>::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<ITYPE, OTYPE>::TreeEnsembleCommon(int parallel_tree, int para
|
|||
}
|
||||
|
||||
template <typename ITYPE, typename OTYPE>
|
||||
void TreeEnsembleCommon<ITYPE, OTYPE>::compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, Tensor* label) const {
|
||||
void TreeEnsembleCommon<ITYPE, OTYPE>::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<ITYPE, OTYPE>::compute(concurrency::ThreadPool* ttp, con
|
|||
|
||||
template <typename ITYPE, typename OTYPE>
|
||||
template <typename AGG>
|
||||
void TreeEnsembleCommon<ITYPE, OTYPE>::ComputeAgg(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, Tensor* label, const AGG& agg) const {
|
||||
void TreeEnsembleCommon<ITYPE, OTYPE>::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<ITYPE>();
|
||||
OTYPE* z_data = Z->template MutableData<OTYPE>();
|
||||
int64_t* label_data = label == NULL ? NULL : label->template MutableData<int64_t>();
|
||||
int64_t* label_data = label == nullptr ? nullptr : label->template MutableData<int64_t>();
|
||||
|
||||
if (n_targets_or_classes_ == 1) {
|
||||
if (N == 1) {
|
||||
ScoreValue<OTYPE> 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<ScoreValue<OTYPE>> scores_t(n_trees_, {0, 0});
|
||||
concurrency::ThreadPool::TryBatchParallelFor(
|
||||
ttp,
|
||||
static_cast<int32_t>(this->n_trees_),
|
||||
[&](ptrdiff_t j) {
|
||||
agg.ProcessTreeNodePrediction1(scores_t[j], *ProcessTreeNodeLeave(this->roots_[j], x_data));
|
||||
SafeInt<int32_t>(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<ITYPE, OTYPE>::ComputeAgg(concurrency::ThreadPool* ttp,
|
|||
|
||||
for (int64_t i = 0; i < N; ++i) {
|
||||
score = {0, 0};
|
||||
for (j = 0; j < static_cast<size_t>(n_trees_); ++j)
|
||||
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 == NULL ? NULL : (label_data + i));
|
||||
label_data == nullptr ? nullptr : (label_data + i));
|
||||
}
|
||||
} else {
|
||||
concurrency::ThreadPool::TryBatchParallelFor(
|
||||
ttp,
|
||||
static_cast<int32_t>(N),
|
||||
[&](ptrdiff_t i) {
|
||||
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)
|
||||
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 == NULL ? NULL : (label_data + i));
|
||||
label_data == nullptr ? nullptr : (label_data + i));
|
||||
},
|
||||
0);
|
||||
}
|
||||
|
|
@ -309,30 +318,32 @@ void TreeEnsembleCommon<ITYPE, OTYPE>::ComputeAgg(concurrency::ThreadPool* ttp,
|
|||
} else {
|
||||
if (N == 1) {
|
||||
std::vector<ScoreValue<OTYPE>> 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<int32_t>(concurrency::ThreadPool::NumThreads(ttp), SafeInt<int32_t>(n_trees_));
|
||||
std::mutex merge_mutex;
|
||||
concurrency::ThreadPool::TrySimpleParallelFor(
|
||||
ttp,
|
||||
static_cast<int32_t>(nth),
|
||||
[&](ptrdiff_t th) {
|
||||
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});
|
||||
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<std::mutex> 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<ScoreValue<OTYPE>> scores(n_targets_or_classes_);
|
||||
|
|
@ -340,35 +351,36 @@ void TreeEnsembleCommon<ITYPE, OTYPE>::ComputeAgg(concurrency::ThreadPool* ttp,
|
|||
|
||||
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)
|
||||
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<int32_t>(concurrency::ThreadPool::NumThreads(ttp), SafeInt<int32_t>(N));
|
||||
concurrency::ThreadPool::TrySimpleParallelFor(
|
||||
ttp,
|
||||
static_cast<int32_t>(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<ScoreValue<OTYPE>> 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<OTYPE>({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<ITYPE, OTYPE>::TreeEnsembleCommonClassifier(
|
|||
const std::vector<int64_t>& class_treeids,
|
||||
const std::vector<OTYPE>& class_weights,
|
||||
const std::vector<std::string>& classlabels_strings,
|
||||
const std::vector<int64_t>& classlabels_int64s) : TreeEnsembleCommon<ITYPE, OTYPE>(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<int64_t>& classlabels_int64s)
|
||||
: TreeEnsembleCommon<ITYPE, OTYPE>(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<ITYPE, OTYPE>::TreeEnsembleCommonClassifier(
|
|||
}
|
||||
|
||||
template <typename ITYPE, typename OTYPE>
|
||||
void TreeEnsembleCommonClassifier<ITYPE, OTYPE>::compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, Tensor* label) const {
|
||||
void TreeEnsembleCommonClassifier<ITYPE, OTYPE>::compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z,
|
||||
Tensor* label) const {
|
||||
if (classlabels_strings_.size() == 0) {
|
||||
this->ComputeAgg(
|
||||
ttp, X, Z, label,
|
||||
|
|
|
|||
|
|
@ -54,6 +54,51 @@ TEST(MLOpTest, TreeEnsembleClassifier) {
|
|||
test.Run();
|
||||
}
|
||||
|
||||
TEST(MLOpTest, TreeEnsembleClassifier_N1) {
|
||||
OpTester test("TreeEnsembleClassifier", 1, onnxruntime::kMLDomain);
|
||||
|
||||
std::vector<int64_t> lefts = {1, -1, 3, -1, -1, 1, -1, 3, 4, -1, -1, -1, 1, 2, -1, 4, -1, -1, -1};
|
||||
std::vector<int64_t> rights = {2, -1, 4, -1, -1, 2, -1, 6, 5, -1, -1, -1, 6, 3, -1, 5, -1, -1, -1};
|
||||
std::vector<int64_t> treeids = {0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 2};
|
||||
std::vector<int64_t> nodeids = {0, 1, 2, 3, 4, 0, 1, 2, 3, 4, 5, 6, 0, 1, 2, 3, 4, 5, 6};
|
||||
std::vector<int64_t> featureids = {2, -2, 0, -2, -2, 0, -2, 2, 1, -2, -2, -2, 0, 2, -2, 1, -2, -2, -2};
|
||||
std::vector<float> 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<std::string> 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<int64_t> class_treeids = {0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2};
|
||||
std::vector<int64_t> class_nodeids = {1, 3, 4, 1, 4, 5, 6, 2, 4, 5, 6};
|
||||
std::vector<int64_t> class_classids = {2, 0, 1, 0, 2, 3, 1, 2, 0, 1, 3};
|
||||
std::vector<float> 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<int64_t> classes = {0, 1, 2, 3};
|
||||
std::vector<float> X = {1.f, 0.0f, 0.4f};
|
||||
std::vector<int64_t> results = {0};
|
||||
std::vector<float> scores{7, 0, 0, 0};
|
||||
std::vector<float> probs = {};
|
||||
std::vector<float> 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<float>("X", {N, 3}, X);
|
||||
test.AddOutput<int64_t>("Y", {N}, results);
|
||||
test.AddOutput<float>("Z", {N, static_cast<int64_t>(classes.size())}, scores);
|
||||
test.Run();
|
||||
}
|
||||
|
||||
TEST(MLOpTest, TreeEnsembleClassifierLabels) {
|
||||
OpTester test("TreeEnsembleClassifier", 1, onnxruntime::kMLDomain);
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue