def test_resume_from_checkpoint(ray_start_4_cpus, tmpdir):
train_dataset = ray.data.from_pandas(train_df)
valid_dataset = ray.data.from_pandas(test_df)
trainer = XGBoostTrainer(
scaling_config=scale_config,
label_column="target",
params=params,
num_boost_round=5,
datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
)
result = trainer.fit()
checkpoint = result.checkpoint
xgb_model, _ = load_checkpoint(checkpoint)
assert get_num_trees(xgb_model) == 5
# Move checkpoint to a different directory.
checkpoint_dict = result.checkpoint.to_dict()
checkpoint = Checkpoint.from_dict(checkpoint_dict)
checkpoint_path = checkpoint.to_directory(tmpdir)
resume_from = Checkpoint.from_directory(checkpoint_path)
trainer = XGBoostTrainer(
scaling_config=scale_config,
label_column="target",
params=params,
num_boost_round=5,
datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
resume_from_checkpoint=resume_from,
)
result = trainer.fit()
checkpoint = result.checkpoint
model, _ = load_checkpoint(checkpoint)
assert get_num_trees(model) == 10
def test_fit(ray_start_4_cpus):
train_dataset = ray.data.from_pandas(train_df)
valid_dataset = ray.data.from_pandas(test_df)
trainer = XGBoostTrainer(
scaling_config=scale_config,
label_column="target",
params=params,
datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
)
trainer.fit()
def test_checkpoint_freq(ray_start_4_cpus, freq_end_expected):
freq, end, expected = freq_end_expected
train_dataset = ray.data.from_pandas(train_df)
valid_dataset = ray.data.from_pandas(test_df)
trainer = XGBoostTrainer(
run_config=ray.air.RunConfig(
checkpoint_config=ray.air.CheckpointConfig(
checkpoint_frequency=freq, checkpoint_at_end=end)),
scaling_config=scale_config,
label_column="target",
params=params,
num_boost_round=25,
datasets={
TRAIN_DATASET_KEY: train_dataset,
"valid": valid_dataset
},
)
result = trainer.fit()
# Assert number of checkpoints
assert len(result.best_checkpoints) == expected, str([
(metrics["training_iteration"], _cp._local_path)
for _cp, metrics in result.best_checkpoints
])
# Assert checkpoint numbers are increasing
cp_paths = [cp._local_path for cp, _ in result.best_checkpoints]
assert cp_paths == sorted(cp_paths), str(cp_paths)
def test_preprocessor_in_checkpoint(ray_start_4_cpus, tmpdir):
train_dataset = ray.data.from_pandas(train_df)
valid_dataset = ray.data.from_pandas(test_df)
class DummyPreprocessor(Preprocessor):
def __init__(self):
super().__init__()
self.is_same = True
def fit(self, dataset):
self.fitted_ = True
def _transform_pandas(self, df: "pd.DataFrame") -> "pd.DataFrame":
return df
trainer = XGBoostTrainer(
scaling_config=scale_config,
label_column="target",
params=params,
datasets={TRAIN_DATASET_KEY: train_dataset, "valid": valid_dataset},
preprocessor=DummyPreprocessor(),
)
result = trainer.fit()
# Move checkpoint to a different directory.
checkpoint_dict = result.checkpoint.to_dict()
checkpoint = Checkpoint.from_dict(checkpoint_dict)
checkpoint_path = checkpoint.to_directory(tmpdir)
resume_from = Checkpoint.from_directory(checkpoint_path)
model, preprocessor = load_checkpoint(resume_from)
assert get_num_trees(model) == 10
assert preprocessor.is_same
assert preprocessor.fitted_
def run_xgboost_training(data_path: str, num_workers: int):
ds = data.read_parquet(data_path)
params = {
"objective": "binary:logistic",
"eval_metric": ["logloss", "error"],
}
trainer = XGBoostTrainer(
scaling_config=ScalingConfig(
num_workers=num_workers,
resources_per_worker={"CPU": 12},
),
label_column="labels",
params=params,
datasets={"train": ds},
)
result = trainer.fit()
checkpoint = XGBoostCheckpoint.from_checkpoint(result.checkpoint)
xgboost_model = checkpoint.get_model()
xgboost_model.save_model(_XGB_MODEL_PATH)
ray.shutdown()
trainer = XGBoostTrainer(
scaling_config={
"num_workers": num_workers,
"use_gpu": use_gpu,
},
label_column="target",
params=params,
datasets={
"train": train_dataset,
"valid": valid_dataset
},
preprocessor=preprocessor,
num_boost_round=20,
)
result = trainer.fit()
print(result.metrics)
# __air_xgb_train_end__
# __air_xgb_batchpred_start__
from ray.train.batch_predictor import BatchPredictor
from ray.train.xgboost import XGBoostPredictor
batch_predictor = BatchPredictor.from_checkpoint(result.checkpoint,
XGBoostPredictor)
predicted_probabilities = batch_predictor.predict(test_dataset)
print("PREDICTED PROBABILITIES")
predicted_probabilities.show()
shap_values = batch_predictor.predict(test_dataset, pred_contribs=True)