From 037dd0820d36c554460a7215bd792c534040ea17 Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Wed, 15 Sep 2021 19:09:04 +0800 Subject: [PATCH] Implement `__sklearn_is_fitted__`. (#7230) --- python-package/xgboost/sklearn.py | 5 ++++- tests/python-gpu/test_gpu_with_sklearn.py | 2 +- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/python-package/xgboost/sklearn.py b/python-package/xgboost/sklearn.py index 999caae45..ed91c4b77 100644 --- a/python-package/xgboost/sklearn.py +++ b/python-package/xgboost/sklearn.py @@ -435,6 +435,9 @@ class XGBModel(XGBModelBase): '''Tags used for scikit-learn data validation.''' return {'allow_nan': True, 'no_validation': True} + def __sklearn_is_fitted__(self) -> bool: + return hasattr(self, "_Booster") + def get_booster(self) -> Booster: """Get the underlying xgboost Booster of this model. @@ -444,7 +447,7 @@ class XGBModel(XGBModelBase): ------- booster : a xgboost booster of underlying model """ - if not hasattr(self, '_Booster'): + if not self.__sklearn_is_fitted__(): from sklearn.exceptions import NotFittedError raise NotFittedError('need to call fit or load_model beforehand') return self._Booster diff --git a/tests/python-gpu/test_gpu_with_sklearn.py b/tests/python-gpu/test_gpu_with_sklearn.py index 7658299b9..47a2f44ca 100644 --- a/tests/python-gpu/test_gpu_with_sklearn.py +++ b/tests/python-gpu/test_gpu_with_sklearn.py @@ -19,7 +19,7 @@ def test_gpu_binary_classification(): from sklearn.datasets import load_digits from sklearn.model_selection import KFold - digits = load_digits(2) + digits = load_digits(n_class=2) y = digits['target'] X = digits['data'] kf = KFold(n_splits=2, shuffle=True, random_state=rng)