treelite_xgboost_model_compile_to_shared_library_inference.py

1import treelite
2import treelite_runtime
3import xgboost as xgb
4import numpy as np
5
6# 1. Train a model using XGBoost
7data = np.random.randn(100, 10)
8label = np.random.randint(0, 2, 100)
9dtrain = xgb.DMatrix(data, label=label)
10bst = xgb.train({'max_depth': 2, 'eta': 1, 'objective': 'binary:logistic'}, dtrain, 10)
11
12# 2. Load the XGBoost model into Treelite
13# Treelite can parse the XGBoost model object directly
14model = treelite.Model.from_xgboost(bst)
15
16# 3. Compile the model into a shared library (DLL)
17# We use the 'gcc' toolchain here; you can also use 'msvc' or 'clang'
18toolchain = 'gcc'
19model.export_lib(toolchain=toolchain, libpath='./mymodel.so', params={'parallel_comp': 4})
20
21# 4. Use the compiled library for prediction
22predictor = treelite_runtime.Predictor('./mymodel.so', nthread=1)
23dmat = treelite_runtime.DMatrix(data)
24out = predictor.predict(dmat)
25
26print(out)