mirror of
https://github.com/deepspeedai/DeepSpeed.git
synced 2025-10-20 15:33:51 +08:00
64 lines
1.7 KiB
Python
Executable File
64 lines
1.7 KiB
Python
Executable File
'''Copyright The Microsoft DeepSpeed Team'''
|
|
|
|
from .utils import *
|
|
|
|
try:
|
|
import xgboost as xgb
|
|
except ImportError:
|
|
xgb = None
|
|
|
|
|
|
class XGBoostCostModel():
|
|
|
|
def __init__(self, loss_type, num_threads=None, log_interval=25, upper_model=None):
|
|
|
|
assert xgb is not None, "missing requirements, please install deepspeed w. 'autotuning_ml' extra."
|
|
|
|
self.loss_type = loss_type
|
|
|
|
if loss_type == "reg":
|
|
self.xgb_params = {
|
|
"max_depth": 3,
|
|
"gamma": 0.0001,
|
|
"min_child_weight": 1,
|
|
"subsample": 1.0,
|
|
"eta": 0.3,
|
|
"lambda": 1.0,
|
|
"alpha": 0,
|
|
"objective": "reg:linear",
|
|
}
|
|
elif loss_type == "rank":
|
|
self.xgb_params = {
|
|
"max_depth": 3,
|
|
"gamma": 0.0001,
|
|
"min_child_weight": 1,
|
|
"subsample": 1.0,
|
|
"eta": 0.3,
|
|
"lambda": 1.0,
|
|
"alpha": 0,
|
|
"objective": "rank:pairwise",
|
|
}
|
|
else:
|
|
raise RuntimeError("Invalid loss type: " + loss_type)
|
|
|
|
self.xgb_params["verbosity"] = 0
|
|
if num_threads:
|
|
self.xgb_params["nthread"] = num_threads
|
|
|
|
def fit(self, xs, ys):
|
|
x_train = np.array(xs, dtype=np.float32)
|
|
y_train = np.array(ys, dtype=np.float32)
|
|
y_max = np.max(y_train)
|
|
y_train = y_train / max(y_max, 1e-9)
|
|
|
|
index = np.random.permutation(len(x_train))
|
|
dtrain = xgb.DMatrix(x_train[index], y_train[index])
|
|
|
|
self.bst = xgb.train(self.xgb_params, dtrain)
|
|
|
|
def predict(self, xs):
|
|
|
|
features = xgb.DMatrix(xs)
|
|
|
|
return self.bst.predict(features)
|