xgboost_ranking.py

import numpy as np
import xgboost as xgb
import time

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

n = 10000000
query_groups = 100000
group_size = n / query_groups
m = 10
X = np.random.randn(n, m)
y = np.random.randn(n)
groups = np.array([group_size] * query_groups, dtype=np.uint32)

start = time.time()
clf = xgb.XGBRanker(tree_method='gpu_hist', objective='rank:map', max_depth=6, eta=0.1)
clf.fit(X, y, groups)
print("Ranking time: {}s".format(time.time() - start))