195 lines
7.7 KiB
Python
195 lines
7.7 KiB
Python
import numpy as np
|
|
import xgboost
|
|
import os
|
|
import itertools
|
|
import shutil
|
|
import urllib.request
|
|
import zipfile
|
|
import sys
|
|
import pytest
|
|
sys.path.append("tests/python")
|
|
|
|
import testing as tm # noqa
|
|
|
|
pytestmark = pytest.mark.timeout(10)
|
|
|
|
class TestRanking:
|
|
@classmethod
|
|
def setup_class(cls):
|
|
"""
|
|
Download and setup the test fixtures
|
|
"""
|
|
from sklearn.datasets import load_svmlight_files
|
|
# download the test data
|
|
cls.dpath = os.path.join(tm.PROJECT_ROOT, "demo/rank/")
|
|
src = 'https://s3-us-west-2.amazonaws.com/xgboost-examples/MQ2008.zip'
|
|
target = os.path.join(cls.dpath, "MQ2008.zip")
|
|
|
|
if os.path.exists(cls.dpath) and os.path.exists(target):
|
|
print("Skipping dataset download...")
|
|
else:
|
|
urllib.request.urlretrieve(url=src, filename=target)
|
|
with zipfile.ZipFile(target, 'r') as f:
|
|
f.extractall(path=cls.dpath)
|
|
|
|
(x_train, y_train, qid_train, x_test, y_test, qid_test,
|
|
x_valid, y_valid, qid_valid) = load_svmlight_files(
|
|
(cls.dpath + "MQ2008/Fold1/train.txt",
|
|
cls.dpath + "MQ2008/Fold1/test.txt",
|
|
cls.dpath + "MQ2008/Fold1/vali.txt"),
|
|
query_id=True, zero_based=False)
|
|
# instantiate the matrices
|
|
cls.dtrain = xgboost.DMatrix(x_train, y_train)
|
|
cls.dvalid = xgboost.DMatrix(x_valid, y_valid)
|
|
cls.dtest = xgboost.DMatrix(x_test, y_test)
|
|
# set the group counts from the query IDs
|
|
cls.dtrain.set_group([len(list(items))
|
|
for _key, items in itertools.groupby(qid_train)])
|
|
cls.dtest.set_group([len(list(items))
|
|
for _key, items in itertools.groupby(qid_test)])
|
|
cls.dvalid.set_group([len(list(items))
|
|
for _key, items in itertools.groupby(qid_valid)])
|
|
# save the query IDs for testing
|
|
cls.qid_train = qid_train
|
|
cls.qid_test = qid_test
|
|
cls.qid_valid = qid_valid
|
|
|
|
def setup_weighted(x, y, groups):
|
|
# Setup weighted data
|
|
data = xgboost.DMatrix(x, y)
|
|
groups_segment = [len(list(items))
|
|
for _key, items in itertools.groupby(groups)]
|
|
data.set_group(groups_segment)
|
|
n_groups = len(groups_segment)
|
|
weights = np.ones((n_groups,))
|
|
data.set_weight(weights)
|
|
return data
|
|
|
|
cls.dtrain_w = setup_weighted(x_train, y_train, qid_train)
|
|
cls.dtest_w = setup_weighted(x_test, y_test, qid_test)
|
|
cls.dvalid_w = setup_weighted(x_valid, y_valid, qid_valid)
|
|
|
|
# model training parameters
|
|
cls.params = {'booster': 'gbtree',
|
|
'tree_method': 'gpu_hist',
|
|
'gpu_id': 0,
|
|
'predictor': 'gpu_predictor'}
|
|
cls.cpu_params = {'booster': 'gbtree',
|
|
'tree_method': 'hist',
|
|
'gpu_id': -1,
|
|
'predictor': 'cpu_predictor'}
|
|
|
|
@classmethod
|
|
def teardown_class(cls):
|
|
"""
|
|
Cleanup test artifacts from download and unpacking
|
|
:return:
|
|
"""
|
|
os.remove(os.path.join(cls.dpath, "MQ2008.zip"))
|
|
shutil.rmtree(os.path.join(cls.dpath, "MQ2008"))
|
|
|
|
@classmethod
|
|
def __test_training_with_rank_objective(cls, rank_objective, metric_name, tolerance=1e-02):
|
|
"""
|
|
Internal method that trains the dataset using the rank objective on GPU and CPU, evaluates
|
|
the metric and determines if the delta between the metric is within the tolerance level
|
|
:return:
|
|
"""
|
|
# specify validations set to watch performance
|
|
watchlist = [(cls.dtest, 'eval'), (cls.dtrain, 'train')]
|
|
|
|
num_trees = 100
|
|
check_metric_improvement_rounds = 10
|
|
|
|
evals_result = {}
|
|
cls.params['objective'] = rank_objective
|
|
cls.params['eval_metric'] = metric_name
|
|
bst = xgboost.train(
|
|
cls.params, cls.dtrain, num_boost_round=num_trees,
|
|
early_stopping_rounds=check_metric_improvement_rounds,
|
|
evals=watchlist, evals_result=evals_result)
|
|
gpu_map_metric = evals_result['train'][metric_name][-1]
|
|
|
|
evals_result = {}
|
|
cls.cpu_params['objective'] = rank_objective
|
|
cls.cpu_params['eval_metric'] = metric_name
|
|
bstc = xgboost.train(
|
|
cls.cpu_params, cls.dtrain, num_boost_round=num_trees,
|
|
early_stopping_rounds=check_metric_improvement_rounds,
|
|
evals=watchlist, evals_result=evals_result)
|
|
cpu_map_metric = evals_result['train'][metric_name][-1]
|
|
|
|
assert np.allclose(gpu_map_metric, cpu_map_metric, tolerance,
|
|
tolerance)
|
|
assert np.allclose(bst.best_score, bstc.best_score, tolerance,
|
|
tolerance)
|
|
|
|
evals_result_weighted = {}
|
|
watchlist = [(cls.dtest_w, 'eval'), (cls.dtrain_w, 'train')]
|
|
bst_w = xgboost.train(
|
|
cls.params, cls.dtrain_w, num_boost_round=num_trees,
|
|
early_stopping_rounds=check_metric_improvement_rounds,
|
|
evals=watchlist, evals_result=evals_result_weighted)
|
|
weighted_metric = evals_result_weighted['train'][metric_name][-1]
|
|
# GPU Ranking is not deterministic due to `AtomicAddGpair`,
|
|
# remove tolerance once the issue is resolved.
|
|
# https://github.com/dmlc/xgboost/issues/5561
|
|
assert np.allclose(bst_w.best_score, bst.best_score,
|
|
tolerance, tolerance)
|
|
assert np.allclose(weighted_metric, gpu_map_metric,
|
|
tolerance, tolerance)
|
|
|
|
def test_training_rank_pairwise_map_metric(self):
|
|
"""
|
|
Train an XGBoost ranking model with pairwise objective function and compare map metric
|
|
"""
|
|
self.__test_training_with_rank_objective('rank:pairwise', 'map')
|
|
|
|
def test_training_rank_pairwise_auc_metric(self):
|
|
"""
|
|
Train an XGBoost ranking model with pairwise objective function and compare auc metric
|
|
"""
|
|
self.__test_training_with_rank_objective('rank:pairwise', 'auc')
|
|
|
|
def test_training_rank_pairwise_ndcg_metric(self):
|
|
"""
|
|
Train an XGBoost ranking model with pairwise objective function and compare ndcg metric
|
|
"""
|
|
self.__test_training_with_rank_objective('rank:pairwise', 'ndcg')
|
|
|
|
def test_training_rank_ndcg_map(self):
|
|
"""
|
|
Train an XGBoost ranking model with ndcg objective function and compare map metric
|
|
"""
|
|
self.__test_training_with_rank_objective('rank:ndcg', 'map')
|
|
|
|
def test_training_rank_ndcg_auc(self):
|
|
"""
|
|
Train an XGBoost ranking model with ndcg objective function and compare auc metric
|
|
"""
|
|
self.__test_training_with_rank_objective('rank:ndcg', 'auc')
|
|
|
|
def test_training_rank_ndcg_ndcg(self):
|
|
"""
|
|
Train an XGBoost ranking model with ndcg objective function and compare ndcg metric
|
|
"""
|
|
self.__test_training_with_rank_objective('rank:ndcg', 'ndcg')
|
|
|
|
def test_training_rank_map_map(self):
|
|
"""
|
|
Train an XGBoost ranking model with map objective function and compare map metric
|
|
"""
|
|
self.__test_training_with_rank_objective('rank:map', 'map')
|
|
|
|
def test_training_rank_map_auc(self):
|
|
"""
|
|
Train an XGBoost ranking model with map objective function and compare auc metric
|
|
"""
|
|
self.__test_training_with_rank_objective('rank:map', 'auc')
|
|
|
|
def test_training_rank_map_ndcg(self):
|
|
"""
|
|
Train an XGBoost ranking model with map objective function and compare ndcg metric
|
|
"""
|
|
self.__test_training_with_rank_objective('rank:map', 'ndcg')
|