Function test_sklearn_model

tests/python/test_model_io.py:281–337

(tmp_path: Path)

@pytest.mark.skipif(**tm.no_sklearn())
def test_sklearn_model(tmp_path: Path) -> None:
    from sklearn.datasets import load_digits
    from sklearn.model_selection import train_test_split

    model_path = tmp_path / "digits.model.json"
    save_load_model(str(model_path))

    model_path = tmp_path / "digits.model.ubj"
    digits = load_digits(n_class=2)
    y = digits["target"]
    X = digits["data"]
    booster = xgb.train(
        {"tree_method": "hist", "objective": "binary:logistic"},
        dtrain=xgb.DMatrix(X, y),
        num_boost_round=4,
    )
    predt_0 = booster.predict(xgb.DMatrix(X))
    booster.save_model(model_path)
    cls = xgb.XGBClassifier()
    cls.load_model(model_path)

    proba = cls.predict_proba(X)
    assert proba.shape[0] == X.shape[0]
    assert proba.shape[1] == 2  # binary

    predt_1 = cls.predict_proba(X)[:, 1]
    assert np.allclose(predt_0, predt_1)

    cls = xgb.XGBModel()
    cls.load_model(model_path)
    predt_1 = cls.predict(X)
    assert np.allclose(predt_0, predt_1)

    # mclass
    X, y = load_digits(n_class=10, return_X_y=True)
    # small test_size to force early stop
    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.01, random_state=1
    )
    clf = xgb.XGBClassifier(
        n_estimators=64, tree_method="hist", early_stopping_rounds=2
    )
    clf.fit(X_train, y_train, eval_set=[(X_test, y_test)])
    score = clf.best_score
    intercept = clf.intercept_
    clf.save_model(model_path)

    clf = xgb.XGBClassifier()
    clf.load_model(model_path)
    assert clf.classes_.size == 10
    assert clf.objective == "multi:softprob"
    np.testing.assert_allclose(intercept, clf.intercept_)

    np.testing.assert_equal(clf.classes_, np.arange(10))
    assert clf.n_classes_ == 10

    assert clf.best_score == score