| 372 | @pytest.mark.skipif(**tm.no_sklearn()) |
| 373 | @pytest.mark.parametrize("booster_name", ["gbtree", "dart"]) |
| 374 | def test_slice(self, booster_name: str) -> None: |
| 375 | from sklearn.datasets import make_classification |
| 376 | |
| 377 | num_classes = 3 |
| 378 | X, y = make_classification( |
| 379 | n_samples=1000, n_informative=5, n_classes=num_classes |
| 380 | ) |
| 381 | dtrain = xgb.DMatrix(data=X, label=y) |
| 382 | num_parallel_tree = 4 |
| 383 | num_boost_round = 16 |
| 384 | total_trees = num_parallel_tree * num_classes * num_boost_round |
| 385 | booster = xgb.train( |
| 386 | { |
| 387 | "num_parallel_tree": num_parallel_tree, |
| 388 | "subsample": 0.5, |
| 389 | "num_class": num_classes, |
| 390 | "booster": booster_name, |
| 391 | "objective": "multi:softprob", |
| 392 | }, |
| 393 | num_boost_round=num_boost_round, |
| 394 | dtrain=dtrain, |
| 395 | ) |
| 396 | booster.feature_types = ["q"] * X.shape[1] |
| 397 | |
| 398 | assert len(booster.get_dump()) == total_trees |
| 399 | |
| 400 | assert booster[...].num_boosted_rounds() == num_boost_round |
| 401 | |
| 402 | self.run_slice( |
| 403 | booster, dtrain, num_parallel_tree, num_classes, num_boost_round, False |
| 404 | ) |
| 405 | |
| 406 | bytesarray = booster.save_raw(raw_format="ubj") |
| 407 | booster = xgb.Booster(model_file=bytesarray) |
| 408 | self.run_slice( |
| 409 | booster, dtrain, num_parallel_tree, num_classes, num_boost_round, False |
| 410 | ) |
| 411 | |
| 412 | @pytest.mark.skipif(**tm.no_pandas()) |
| 413 | @pytest.mark.parametrize("ext", ["json", "ubj"]) |