Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add special handling for multi:softmax in sklearn predict (#7607)

Browse Source 
* Add special handling for multi:softmax in sklearn predict

* Add test coverage
This commit is contained in:
Philip Hyunsu Cho
2022-01-29 15:54:49 -08:00
committed by GitHub
parent 7f738e7f6f
commit b4340abf56
2 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
tests/python/test_with_sklearn.py
View File

@@ -36,7 +36,8 @@ def test_binary_classification():
assert err < 0.1
def test_multiclass_classification():
@pytest.mark.parametrize('objective', ['multi:softmax', 'multi:softprob'])
def test_multiclass_classification(objective):
from sklearn.datasets import load_iris
from sklearn.model_selection import KFold
@@ -54,7 +55,7 @@ def test_multiclass_classification():
X = iris['data']
kf = KFold(n_splits=2, shuffle=True, random_state=rng)
for train_index, test_index in kf.split(X, y):
xgb_model = xgb.XGBClassifier().fit(X[train_index], y[train_index])
xgb_model = xgb.XGBClassifier(objective=objective).fit(X[train_index], y[train_index])
assert (xgb_model.get_booster().num_boosted_rounds() ==
xgb_model.n_estimators)
preds = xgb_model.predict(X[test_index])