Avoid OMP reduction in AUC. (#7362)

This commit is contained in:
Jiaming Yuan 2021-10-28 05:03:52 +08:00 committed by GitHub
parent ac9bfaa4f2
commit d05754f558
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -7,6 +7,7 @@
#include <functional>
#include <limits>
#include <memory>
#include <numeric>
#include <utility>
#include <tuple>
#include <vector>
@ -203,17 +204,15 @@ std::pair<float, uint32_t> RankingAUC(std::vector<float> const &predts,
MetaInfo const &info, int32_t n_threads) {
CHECK_GE(info.group_ptr_.size(), 2);
uint32_t n_groups = info.group_ptr_.size() - 1;
float sum_auc = 0;
auto s_predts = common::Span<float const>{predts};
auto s_labels = info.labels_.ConstHostSpan();
auto s_weights = info.weights_.ConstHostSpan();
std::atomic<uint32_t> invalid_groups{0};
dmlc::OMPException omp_handler;
#pragma omp parallel for reduction(+:sum_auc) num_threads(n_threads)
for (omp_ulong g = 1; g < info.group_ptr_.size(); ++g) {
omp_handler.Run([&]() {
std::vector<double> auc_tloc(n_threads, 0);
common::ParallelFor(n_groups, n_threads, [&](size_t g) {
g += 1; // indexing needs to start from 1
size_t cnt = info.group_ptr_[g] - info.group_ptr_[g - 1];
float w = s_weights.empty() ? 1.0f : s_weights[g - 1];
auto g_predts = s_predts.subspan(info.group_ptr_[g - 1], cnt);
@ -235,10 +234,9 @@ std::pair<float, uint32_t> RankingAUC(std::vector<float> const &predts,
auc = 0;
}
}
sum_auc += auc;
auc_tloc[omp_get_thread_num()] += auc;
});
}
omp_handler.Rethrow();
float sum_auc = std::accumulate(auc_tloc.cbegin(), auc_tloc.cend(), 0.0);
return std::make_pair(sum_auc, n_groups - invalid_groups);
}