dmitryikh / leaves

pure Go implementation of prediction part for GBRT (Gradient Boosting Regression Trees) models from popular frameworks
MIT License
432 stars 72 forks source link

Support outputting leaf indices for all the `predict*` functions #75

Closed fredrikluo closed 4 years ago

fredrikluo commented 4 years ago

Thanks to the popular paper https://research.fb.com/wp-content/uploads/2016/11/practical-lessons-from-predicting-clicks-on-ads-at-facebook.pdf

Many people use GBDT to extract features from a dataset and then instead of predicting results directly. The extracted features are the leaf indices from each estimator which makes the decision.

This is achieved by setting predleaf with lightGBM https://lightgbm.readthedocs.io/en/latest/pythonapi/lightgbm.Booster.html#lightgbm.Booster.predict

I am adding the same support in this pull request. Basically, the user can set this parameter and then predict function will return the leaf indices.

The test data are generated by lightGBM to make sure that the indices are generated in the same way.