saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-06 04:28:32 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Fixes #4688, remove CPUAllocator in TreeEnsemble (#5375)

Browse source
This commit is contained in:
Xavier Dupré 2020-10-09 11:26:07 +02:00 committed by GitHub
parent d4507e9331
commit 621fdb44e5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 15 additions and 14 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

@ -173,7 +173,7 @@ common::Status TreeEnsembleClassifier<T>::Compute(OpKernelContext* context) cons
Tensor* Y = context->Output(0, {N});
Tensor* Z = context->Output(1, {N, tree_ensemble_.get_class_count()});
tree_ensemble_.compute(context->GetOperatorThreadPool(), &X, Z, Y);
tree_ensemble_.compute(context, &X, Z, Y);
return Status::OK();
}

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

@ -52,7 +52,7 @@ class TreeEnsembleCommon {
const std::vector<int64_t>& target_class_treeids,
const std::vector<OTYPE>& target_class_weights);
void compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, Tensor* label) const;
void compute(OpKernelContext* ctx, const Tensor* X, Tensor* Z, Tensor* label) const;
protected:
TreeNodeElement<OTYPE>* ProcessTreeNodeLeave(
@ -216,33 +216,33 @@ 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,
void TreeEnsembleCommon<ITYPE, OTYPE>::compute(OpKernelContext* ctx, const Tensor* X, Tensor* Z,
Tensor* label) const {
switch (aggregate_function_) {
case AGGREGATE_FUNCTION::AVERAGE:
ComputeAgg(
ttp, X, Z, label,
ctx->GetOperatorThreadPool(), X, Z, label,
TreeAggregatorAverage<ITYPE, OTYPE>(
roots_.size(), n_targets_or_classes_,
post_transform_, base_values_));
return;
case AGGREGATE_FUNCTION::SUM:
ComputeAgg(
ttp, X, Z, label,
ctx->GetOperatorThreadPool(), X, Z, label,
TreeAggregatorSum<ITYPE, OTYPE>(
roots_.size(), n_targets_or_classes_,
post_transform_, base_values_));
return;
case AGGREGATE_FUNCTION::MIN:
ComputeAgg(
ttp, X, Z, label,
ctx->GetOperatorThreadPool(), X, Z, label,
TreeAggregatorMin<ITYPE, OTYPE>(
roots_.size(), n_targets_or_classes_,
post_transform_, base_values_));
return;
case AGGREGATE_FUNCTION::MAX:
ComputeAgg(
ttp, X, Z, label,
ctx->GetOperatorThreadPool(), X, Z, label,
TreeAggregatorMax<ITYPE, OTYPE>(
roots_.size(), n_targets_or_classes_,
post_transform_, base_values_));
@ -525,7 +525,7 @@ class TreeEnsembleCommonClassifier : TreeEnsembleCommon<ITYPE, OTYPE> {
int64_t get_class_count() const { return this->n_targets_or_classes_; }
void compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z, Tensor* label) const;
void compute(OpKernelContext* ctx, const Tensor* X, Tensor* Z, Tensor* label) const;
};
template <typename ITYPE, typename OTYPE>
@ -589,11 +589,11 @@ TreeEnsembleCommonClassifier<ITYPE, OTYPE>::TreeEnsembleCommonClassifier(
}
template <typename ITYPE, typename OTYPE>
void TreeEnsembleCommonClassifier<ITYPE, OTYPE>::compute(concurrency::ThreadPool* ttp, const Tensor* X, Tensor* Z,
void TreeEnsembleCommonClassifier<ITYPE, OTYPE>::compute(OpKernelContext* ctx, const Tensor* X, Tensor* Z,
Tensor* label) const {
if (classlabels_strings_.size() == 0) {
this->ComputeAgg(
ttp, X, Z, label,
ctx->GetOperatorThreadPool(), X, Z, label,
TreeAggregatorClassifier<ITYPE, OTYPE>(
this->roots_.size(), this->n_targets_or_classes_,
this->post_transform_, this->base_values_,
@ -601,10 +601,11 @@ void TreeEnsembleCommonClassifier<ITYPE, OTYPE>::compute(concurrency::ThreadPool
weights_are_all_positive_));
} else {
int64_t N = X->Shape().NumDimensions() == 1 ? 1 : X->Shape()[0];
std::shared_ptr<IAllocator> allocator = std::make_shared<CPUAllocator>();
Tensor label_int64(DataTypeImpl::GetType<int64_t>(), TensorShape({N}), allocator);
AllocatorPtr alloc;
ORT_THROW_IF_ERROR(ctx->GetTempSpaceAllocator(&alloc));
Tensor label_int64(DataTypeImpl::GetType<int64_t>(), TensorShape({N}), alloc);
this->ComputeAgg(
ttp, X, Z, &label_int64,
ctx->GetOperatorThreadPool(), X, Z, &label_int64,
TreeAggregatorClassifier<ITYPE, OTYPE>(
this->roots_.size(), this->n_targets_or_classes_,
this->post_transform_, this->base_values_,

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

@ -56,7 +56,7 @@ common::Status TreeEnsembleRegressor<T>::Compute(OpKernelContext* context) const
int64_t N = X->Shape().NumDimensions() == 1 ? 1 : X->Shape()[0];
Tensor* Y = context->Output(0, {N, tree_ensemble_.n_targets_or_classes_});
tree_ensemble_.compute(context->GetOperatorThreadPool(), X, Y, NULL);
tree_ensemble_.compute(context, X, Y, NULL);
return Status::OK();
}