Unify max nodes. (#5497)

This commit is contained in:
Jiaming Yuan 2020-04-10 19:26:35 +08:00 committed by GitHub
parent bd653fad4c
commit 866a477319
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 16 additions and 22 deletions

View File

@ -53,16 +53,7 @@ void FeatureInteractionConstraint::Configure(
}
// --- Initialize allowed features attached to nodes.
if (param.max_depth == 0 && param.max_leaves == 0) {
LOG(FATAL) << "Max leaves and max depth cannot both be unconstrained for gpu_hist.";
}
int32_t n_nodes {0};
if (param.max_depth != 0) {
n_nodes = std::pow(2, param.max_depth + 1);
} else {
n_nodes = param.max_leaves * 2 - 1;
}
CHECK_NE(n_nodes, 0);
int32_t n_nodes { param.MaxNodes() };
node_constraints_.resize(n_nodes);
node_constraints_storage_.resize(n_nodes);
for (auto& n : node_constraints_storage_) {

View File

@ -230,6 +230,20 @@ struct TrainParam : public XGBoostParameter<TrainParam> {
CHECK_GT(ret, 0U);
return ret;
}
bst_node_t MaxNodes() const {
if (this->max_depth == 0 && this->max_leaves == 0) {
LOG(FATAL) << "Max leaves and max depth cannot both be unconstrained.";
}
bst_node_t n_nodes{0};
if (this->max_leaves > 0) {
n_nodes = this->max_leaves * 2 - 1;
} else {
n_nodes = (1 << (this->max_depth + 1)) - 1;
}
CHECK_NE(n_nodes, 0);
return n_nodes;
}
};
/*! \brief Loss functions */

View File

@ -178,10 +178,5 @@ struct SumCallbackOp {
return old_prefix;
}
};
// Total number of nodes in tree, given depth
XGBOOST_DEVICE inline int MaxNodesDepth(int depth) {
return (1 << (depth + 1)) - 1;
}
} // namespace tree
} // namespace xgboost

View File

@ -927,13 +927,7 @@ struct GPUHistMakerDevice {
template <typename GradientSumT>
inline void GPUHistMakerDevice<GradientSumT>::InitHistogram() {
CHECK(!(param.max_leaves == 0 && param.max_depth == 0))
<< "Max leaves and max depth cannot both be unconstrained for "
"gpu_hist.";
int max_nodes =
param.max_leaves > 0 ? param.max_leaves * 2 : MaxNodesDepth(param.max_depth);
bst_node_t max_nodes { param.MaxNodes() };
ba.Allocate(device_id,
&prediction_cache, n_rows,
&node_sum_gradients_d, max_nodes,