From 801116c30707d818e084cf95201db5f0ec361b17 Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Sun, 13 Aug 2023 23:41:49 +0800 Subject: [PATCH] Test scikit-learn model IO with gblinear. (#9459) --- tests/python/test_with_sklearn.py | 31 +++++++++++++++++++++++++------ 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/tests/python/test_with_sklearn.py b/tests/python/test_with_sklearn.py index 9a58b7277..69f144caf 100644 --- a/tests/python/test_with_sklearn.py +++ b/tests/python/test_with_sklearn.py @@ -792,19 +792,19 @@ def test_kwargs_grid_search(): from sklearn import datasets from sklearn.model_selection import GridSearchCV - params = {'tree_method': 'hist'} - clf = xgb.XGBClassifier(n_estimators=1, learning_rate=1.0, **params) - assert clf.get_params()['tree_method'] == 'hist' - # 'max_leaves' is not a default argument of XGBClassifier + params = {"tree_method": "hist"} + clf = xgb.XGBClassifier(n_estimators=3, **params) + assert clf.get_params()["tree_method"] == "hist" + # 'eta' is not a default argument of XGBClassifier # Check we can still do grid search over this parameter - search_params = {'max_leaves': range(2, 5)} + search_params = {"eta": [0, 0.2, 0.4]} grid_cv = GridSearchCV(clf, search_params, cv=5) iris = datasets.load_iris() grid_cv.fit(iris.data, iris.target) # Expect unique results for each parameter value # This confirms sklearn is able to successfully update the parameter - means = grid_cv.cv_results_['mean_test_score'] + means = grid_cv.cv_results_["mean_test_score"] assert len(means) == len(set(means)) @@ -928,6 +928,25 @@ def save_load_model(model_path): xgb_model = xgb.XGBModel() xgb_model.load_model(model_path) + clf = xgb.XGBClassifier(booster="gblinear", early_stopping_rounds=1) + clf.fit(X, y, eval_set=[(X, y)]) + best_iteration = clf.best_iteration + best_score = clf.best_score + predt_0 = clf.predict(X) + clf.save_model(model_path) + clf.load_model(model_path) + predt_1 = clf.predict(X) + np.testing.assert_allclose(predt_0, predt_1) + assert clf.best_iteration == best_iteration + assert clf.best_score == best_score + + clfpkl = pickle.dumps(clf) + clf = pickle.loads(clfpkl) + predt_2 = clf.predict(X) + np.testing.assert_allclose(predt_0, predt_2) + assert clf.best_iteration == best_iteration + assert clf.best_score == best_score + def test_save_load_model(): with tempfile.TemporaryDirectory() as tempdir: