| 344 | } |
| 345 | |
| 346 | void BoostedTreesEnsembleResource::PostPruneTree(const int32 current_tree, |
| 347 | const int32 logits_dimension) { |
| 348 | // No-op if tree is empty. |
| 349 | auto* tree = tree_ensemble_->mutable_trees(current_tree); |
| 350 | int32 num_nodes = tree->nodes_size(); |
| 351 | if (num_nodes == 0) { |
| 352 | return; |
| 353 | } |
| 354 | |
| 355 | std::vector<int32> nodes_to_delete; |
| 356 | // If a node was pruned, we need to save the change of the prediction from |
| 357 | // this node to its parent, as well as the parent id. |
| 358 | std::vector<std::pair<int32, std::vector<float>>> nodes_changes; |
| 359 | nodes_changes.reserve(num_nodes); |
| 360 | for (int32 i = 0; i < num_nodes; ++i) { |
| 361 | std::vector<float> prune_logit_changes; |
| 362 | nodes_changes.emplace_back(i, prune_logit_changes); |
| 363 | } |
| 364 | // Prune the tree recursively starting from the root. Each node that has |
| 365 | // negative gain and only leaf children will be pruned recursively up from |
| 366 | // the bottom of the tree. This method returns the list of nodes pruned, and |
| 367 | // updates the nodes in the tree not to refer to those pruned nodes. |
| 368 | RecursivelyDoPostPrunePreparation(current_tree, 0, &nodes_to_delete, |
| 369 | &nodes_changes); |
| 370 | |
| 371 | if (nodes_to_delete.empty()) { |
| 372 | // No pruning happened, and no post-processing needed. |
| 373 | return; |
| 374 | } |
| 375 | |
| 376 | // Sort node ids so they are in asc order. |
| 377 | std::sort(nodes_to_delete.begin(), nodes_to_delete.end()); |
| 378 | |
| 379 | // We need to |
| 380 | // - update split left and right children ids with new indices |
| 381 | // - actually remove the nodes that need to be removed |
| 382 | // - save the information about pruned node so we could recover the |
| 383 | // predictions from cache. Build a map for old node index=>new node index. |
| 384 | // nodes_to_delete contains nodes who's indices should be skipped, in |
| 385 | // ascending order. Save the information about new indices into meta. |
| 386 | std::map<int32, int32> old_to_new_ids; |
| 387 | int32 new_index = 0; |
| 388 | int32 index_for_deleted = 0; |
| 389 | auto* post_prune_meta = tree_ensemble_->mutable_tree_metadata(current_tree) |
| 390 | ->mutable_post_pruned_nodes_meta(); |
| 391 | |
| 392 | for (int32 i = 0; i < num_nodes; ++i) { |
| 393 | if (index_for_deleted < nodes_to_delete.size() && |
| 394 | i == nodes_to_delete[index_for_deleted]) { |
| 395 | // Node i will get removed, |
| 396 | ++index_for_deleted; |
| 397 | // Update meta info that will allow us to use cached predictions from |
| 398 | // those nodes. |
| 399 | int32 new_id; |
| 400 | std::vector<float> logit_changes; |
| 401 | logit_changes.reserve(logits_dimension); |
| 402 | CalculateParentAndLogitUpdate(i, nodes_changes, &new_id, &logit_changes); |
| 403 | auto* meta = post_prune_meta->Add(); |