Parallel model dump for trees. (#7040)

This commit is contained in:
Jiaming Yuan 2021-06-15 14:08:26 +08:00 committed by GitHub
parent 2567404ab6
commit 5c2d7a18c9
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -17,6 +17,8 @@
#include <string> #include <string>
#include <vector> #include <vector>
#include "../common/threading_utils.h"
namespace xgboost { namespace xgboost {
class Json; class Json;
@ -107,12 +109,12 @@ struct GBTreeModel : public Model {
void SaveModel(Json* p_out) const override; void SaveModel(Json* p_out) const override;
void LoadModel(Json const& p_out) override; void LoadModel(Json const& p_out) override;
std::vector<std::string> DumpModel(const FeatureMap& fmap, bool with_stats, std::vector<std::string> DumpModel(const FeatureMap &fmap, bool with_stats,
std::string format) const { std::string format) const {
std::vector<std::string> dump; std::vector<std::string> dump(trees.size());
for (const auto & tree : trees) { common::ParallelFor(static_cast<omp_ulong>(trees.size()), [&](size_t i) {
dump.push_back(tree->DumpModel(fmap, with_stats, format)); dump[i] = trees[i]->DumpModel(fmap, with_stats, format);
} });
return dump; return dump;
} }
void CommitModel(std::vector<std::unique_ptr<RegTree> >&& new_trees, void CommitModel(std::vector<std::unique_ptr<RegTree> >&& new_trees,