Test predict_leaf / predict_per_tree with XGBoost binary classifier
(predict_kind, dataset, num_boost_round)
| 75 | ) |
| 76 | @settings(**standard_settings()) |
| 77 | def test_predict_special_with_binary_classifier(predict_kind, dataset, num_boost_round): |
| 78 | # pylint: disable=too-many-locals |
| 79 | """Test predict_leaf / predict_per_tree with XGBoost binary classifier""" |
| 80 | X, y = dataset |
| 81 | |
| 82 | dtrain = xgb.DMatrix(X, label=y) |
| 83 | param = { |
| 84 | "objective": "binary:logistic", |
| 85 | "base_score": 0.5, |
| 86 | "learning_rate": 0.1, |
| 87 | "max_depth": 6, |
| 88 | } |
| 89 | xgb_model = xgb.train( |
| 90 | param, |
| 91 | dtrain=dtrain, |
| 92 | num_boost_round=num_boost_round, |
| 93 | ) |
| 94 | model: treelite.Model = treelite.frontend.from_xgboost(xgb_model) |
| 95 | assert model.num_tree == num_boost_round |
| 96 | |
| 97 | if predict_kind == "leaf_id": |
| 98 | leaf_pred = treelite.gtil.predict_leaf(model, X) |
| 99 | assert leaf_pred.shape == (X.shape[0], num_boost_round) |
| 100 | xgb_leaf_pred = xgb_model.predict(xgb.DMatrix(X), pred_leaf=True) |
| 101 | assert np.array_equal(leaf_pred, xgb_leaf_pred) |
| 102 | else: |
| 103 | pred_per_tree = treelite.gtil.predict_per_tree(model, X) |
| 104 | assert pred_per_tree.shape == (X.shape[0], num_boost_round, 1) |
| 105 | pred = xgb_model.predict(xgb.DMatrix(X), output_margin=True) |
| 106 | np.testing.assert_almost_equal( |
| 107 | np.sum(pred_per_tree, axis=1).flatten(), pred, decimal=3 |
| 108 | ) |
| 109 | |
| 110 | |
| 111 | @given( |
nothing calls this directly
no test coverage detected