Fix quantile regression demo. (#9991)

This commit is contained in:
Jiaming Yuan 2024-01-17 13:19:08 +08:00 committed by GitHub
parent cacb4b1fdd
commit d07e8b503e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -46,10 +46,11 @@ def quantile_loss(args: argparse.Namespace) -> None:
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=rng) X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=rng)
# We will be using the `hist` tree method, quantile DMatrix can be used to preserve # We will be using the `hist` tree method, quantile DMatrix can be used to preserve
# memory. # memory (which has nothing to do with quantile regression itself, see its document
# for details).
# Do not use the `exact` tree method for quantile regression, otherwise the # Do not use the `exact` tree method for quantile regression, otherwise the
# performance might drop. # performance might drop.
Xy = xgb.QuantileDMatrix(X, y) Xy = xgb.QuantileDMatrix(X_train, y_train)
# use Xy as a reference # use Xy as a reference
Xy_test = xgb.QuantileDMatrix(X_test, y_test, ref=Xy) Xy_test = xgb.QuantileDMatrix(X_test, y_test, ref=Xy)