Run tests for leaf index prediction.
(device: Device, DMatrixT: Type[DMatrix])
| 14 | |
| 15 | # pylint: disable=too-many-locals |
| 16 | def run_predict_leaf(device: Device, DMatrixT: Type[DMatrix]) -> np.ndarray: |
| 17 | """Run tests for leaf index prediction.""" |
| 18 | rows = 100 |
| 19 | cols = 4 |
| 20 | classes = 5 |
| 21 | num_parallel_tree = 4 |
| 22 | num_boost_round = 10 |
| 23 | rng = np.random.RandomState(1994) |
| 24 | X = rng.randn(rows, cols) |
| 25 | y = rng.randint(low=0, high=classes, size=rows) |
| 26 | |
| 27 | m = DMatrixT(X, y) |
| 28 | booster = train( |
| 29 | { |
| 30 | "num_parallel_tree": num_parallel_tree, |
| 31 | "num_class": classes, |
| 32 | "tree_method": "hist", |
| 33 | }, |
| 34 | m, |
| 35 | num_boost_round=num_boost_round, |
| 36 | ) |
| 37 | |
| 38 | booster.set_param({"device": device}) |
| 39 | empty = DMatrixT(np.ones(shape=(0, cols))) |
| 40 | empty_leaf = booster.predict(empty, pred_leaf=True) |
| 41 | assert empty_leaf.shape[0] == 0 |
| 42 | |
| 43 | leaf = booster.predict(m, pred_leaf=True, strict_shape=True) |
| 44 | assert leaf.shape[0] == rows |
| 45 | assert leaf.shape[1] == num_boost_round |
| 46 | assert leaf.shape[2] == classes |
| 47 | assert leaf.shape[3] == num_parallel_tree |
| 48 | |
| 49 | validate_leaf_output(leaf, num_parallel_tree) |
| 50 | |
| 51 | n_iters = np.int32(2) |
| 52 | sliced = booster.predict( |
| 53 | m, |
| 54 | pred_leaf=True, |
| 55 | iteration_range=(0, n_iters), |
| 56 | strict_shape=True, |
| 57 | ) |
| 58 | first = sliced[0, ...] |
| 59 | |
| 60 | assert np.prod(first.shape) == classes * num_parallel_tree * n_iters |
| 61 | |
| 62 | # When there's only 1 tree, the output is a 1 dim vector |
| 63 | booster = train({"tree_method": "hist"}, num_boost_round=1, dtrain=m) |
| 64 | booster.set_param({"device": device}) |
| 65 | assert booster.predict(m, pred_leaf=True).shape == (rows,) |
| 66 | |
| 67 | return leaf |
| 68 | |
| 69 | |
| 70 | def run_base_margin_vs_base_score(device: Device) -> None: |