search

loft-br / xgboost-survival-embeddings

Improving XGBoost survival analysis with embeddings and debiased estimators
https://loft-br.github.io/xgboost-survival-embeddings/
Apache License 2.0
313 stars 51 forks source link

Question: proper way to save/load XGBSEStackedWeibull model #47

Open yerachoi opened 2 years ago

yerachoi commented 2 years ago

Hi. First of all, thank you for developing and sharing XGBSE, which has been of great help for me. I'd like to know how to properly save/load all the key components of a trained XGBSEStackedWeibull model. I have tried out several saving methods but have failed thus far to retrieve the exact same results from the loaded model.

GabrielGimenez commented 2 years ago

We usually use joblib dump and load to save and load models. It's important to fix xgboost, sklearn and lifelines versions when loading, to avoid some difference in the underlying libraries.

so it should be something like

from joblib import dump, load
from xgbse import XGBSEStackedWeibull

...

model = XGBSEStackedWeibull()
model.fit(x_train, y_train)
dump(model, model_path)

...

model = load(model_path)
predictions = model.predict(x_valid)

I'll add a save and load example to the documentation