Basic training continuation.
(tmpdir: str, use_pickle: bool)
| 13 | |
| 14 | |
| 15 | def training_continuation(tmpdir: str, use_pickle: bool) -> None: |
| 16 | """Basic training continuation.""" |
| 17 | # Train 128 iterations in 1 session |
| 18 | X, y = load_breast_cancer(return_X_y=True) |
| 19 | clf = xgboost.XGBClassifier(n_estimators=128, eval_metric="logloss") |
| 20 | clf.fit(X, y, eval_set=[(X, y)]) |
| 21 | print("Total boosted rounds:", clf.get_booster().num_boosted_rounds()) |
| 22 | |
| 23 | # Train 128 iterations in 2 sessions, with the first one runs for 32 iterations and |
| 24 | # the second one runs for 96 iterations |
| 25 | clf = xgboost.XGBClassifier(n_estimators=32, eval_metric="logloss") |
| 26 | clf.fit(X, y, eval_set=[(X, y)]) |
| 27 | assert clf.get_booster().num_boosted_rounds() == 32 |
| 28 | |
| 29 | # load back the model, this could be a checkpoint |
| 30 | if use_pickle: |
| 31 | path = os.path.join(tmpdir, "model-first-32.pkl") |
| 32 | with open(path, "wb") as fd: |
| 33 | pickle.dump(clf, fd) |
| 34 | with open(path, "rb") as fd: |
| 35 | loaded = pickle.load(fd) |
| 36 | else: |
| 37 | path = os.path.join(tmpdir, "model-first-32.json") |
| 38 | clf.save_model(path) |
| 39 | loaded = xgboost.XGBClassifier() |
| 40 | loaded.load_model(path) |
| 41 | |
| 42 | clf = xgboost.XGBClassifier(n_estimators=128 - 32, eval_metric="logloss") |
| 43 | clf.fit(X, y, eval_set=[(X, y)], xgb_model=loaded) |
| 44 | |
| 45 | print("Total boosted rounds:", clf.get_booster().num_boosted_rounds()) |
| 46 | |
| 47 | assert clf.get_booster().num_boosted_rounds() == 128 |
| 48 | |
| 49 | |
| 50 | def training_continuation_early_stop(tmpdir: str, use_pickle: bool) -> None: |
no test coverage detected