xgboost_inplace_prediction.py

import cupy as cp
import GPUtil
import xgboost as xgb
import time

print("Xgboost version: {}".format(xgb.__version__))

n_train = 10000
n_test = 1000
iterations = 20
m = 100
X = cp.random.randn(n_train, m)
X_test = [cp.random.randn(n_test, m) for i in range(iterations)]
y = cp.random.randn(n_train)
dtrain = xgb.DeviceQuantileDMatrix(X, y)
# Train a regression model
bst = xgb.train({
   'tree_method': 'gpu_hist'}, dtrain, 100)
# Standard prediction
start = time.time()
for i in range(iterations):
   pred = bst.predict(xgb.DMatrix(X_test[i]))
print("Standard prediction took: {}".format(time.time() - start))
start = time.time()
# In place prediction
for i in range(iterations):
   pred = bst.inplace_predict(X_test[i])
print("Inplace prediction took: {}".format(time.time() - start))