2021-08-13 20:55:57 +00:00
|
|
|
import xgboost as xgb
|
2021-12-23 12:07:42 +00:00
|
|
|
from sklearn.datasets import load_iris
|
2021-08-13 20:55:57 +00:00
|
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
|
|
|
|
from clearml import Task
|
|
|
|
|
2021-12-23 12:07:42 +00:00
|
|
|
task = Task.init(project_name="examples", task_name="xgboost metric auto reporting")
|
2021-08-13 20:55:57 +00:00
|
|
|
|
2021-12-23 12:07:42 +00:00
|
|
|
X, y = load_iris(return_X_y=True)
|
|
|
|
X_train, X_test, y_train, y_test = train_test_split(
|
|
|
|
X, y, test_size=0.2, random_state=100
|
|
|
|
)
|
2021-08-13 20:55:57 +00:00
|
|
|
|
|
|
|
dtrain = xgb.DMatrix(X_train, label=y_train)
|
|
|
|
dtest = xgb.DMatrix(X_test, label=y_test)
|
|
|
|
|
2021-12-23 12:07:42 +00:00
|
|
|
params = {"objective": "reg:squarederror", "eval_metric": "rmse"}
|
2021-08-13 20:55:57 +00:00
|
|
|
|
|
|
|
bst = xgb.train(
|
2021-12-23 12:07:42 +00:00
|
|
|
params,
|
|
|
|
dtrain,
|
|
|
|
num_boost_round=100,
|
|
|
|
evals=[(dtrain, "train"), (dtest, "test")],
|
|
|
|
verbose_eval=0,
|
2021-08-13 20:55:57 +00:00
|
|
|
)
|
2021-12-23 12:07:42 +00:00
|
|
|
|
|
|
|
bst.save_model("best_model")
|