工作中，很多场景下会用到xgboost模型，如风控、催收、营销、推荐等待。在用xgboost模型进行模型训练的时候，也经常用贝叶斯自动调参来搜索最优的参数。
现在把相关的代码贴出来，供大家参考。
目前是支持了xgboost和lightgbm模型，如果需要支持其他模型，感兴趣的朋友可以自己添加一下。
一、安装包
pip install bayesian-optimization
二、代码示例
GBDTModel类
我写了一个GBDTModel类，用来封装贝叶斯自动调参、模型训练、模型评估、模型保存、参数保存等代码，具体如下：
# -*-coding:utf-8 -*-
import os
import json
import pickle
import numpy as np
import xgboost as xgb
import lightgbm as lgb
from bayes_opt import BayesianOptimization
from sklearn.metrics import mean_absolute_error, roc_auc_score, accuracy_score, confusion_matrix
class GBDTModel(object):
def __init__(self, output_dir, n_thread=10, n_jobs=10, random_state=7):
self.best_model = None
self.best_param = None
self.output_dir = output_dir
os.makedirs(output_dir, exist_ok=True)
self.n_thread = n_thread
self.n_jobs = n_jobs
self.random_state = random_state
self.clf_dict = {
'lgb': lgb.LGBMClassifier,
'xgb': xgb.XGBClassifier
}
self.params_dict = {
'lgb': {
'learning_rate': (0.03, 0.2),
'n_estimators': (100, 1000),
'num_leaves': (20, 100),
'max_depth': (2, 5),
'max_bin': (30, 255),
'min_data_in_bin': (30, 256),
'min_data_in_leaf': (30, 500),
'min_split_gain': (0.1, 10),
'min_child_weight': (0.1, 100),
'min_child_samples': (30, 256),
'colsample_bytree': (0.6, 1.0),
'reg_alpha': (0.3, 10),
'reg_lambda': (0.3, 10),
'subsample': (0.6, 1.0),
'subsample_for_bin': (200000, 400000),
'scale_pos_weight': (1.0, 1.5),
},
'xgb': {
'learning_rate': (0.01, 0.3),
'n_estimators': (100, 800),
'max_depth': (2, 5),
'reg_alpha': (0.00001, 5),
'reg_lambda': (0.00001, 5),
'min_child_weight': (1, 100),
'subsample': (0.6, 1.0),
'colsample_bytree': (0.6, 1.0),
'colsample_bynode': (0.6, 1.0),
'colsample_bylevel': (0.6, 1.0),
'gamma': (0.0, 0.3),
'scale_pos_weight': (1.0, 1.5)
}
}
self.int_params = [
'max_depth', 'n_estimators', 'num_leaves', 'max_bin', 'scale_pos_weight',
'min_child_samples', 'min_data_in_bin', 'min_data_in_leaf', 'subsample_for_bin']
def save_model(self, model_name="model.pkl"):
file = os.path.join(self.output_dir, model_name)
with open(file, 'wb') as f:
pickle.dump(self.best_model, f)
def read_model(self, model_name="model.pkl"):
file = os.path.join(self.output_dir, model_name)
with open(file, 'rb') as f:
self.best_model = pickle.load(f)
def save_best_param(self, file_name="params.json"):
file = os.path.join(self.output_dir, file_name)
with open(file, mode="w", encoding="utf-8") as fw:
json.dump(self.best_param, fw, indent=4, ensure_ascii=False)
def read_best_param(self, file_name="params.json"):
file = os.path.join(self.output_dir, file_name)
with open(file=file, mode="r", encoding="utf-8") as fr:
self.best_param = json.load(fr)
def bayes_optimization(self, train_x, train_y, test_x, test_y, oot_x, oot_y,
clf_name='lgb', target='acc', scale_pos_use=True, **bayes_args):
res_list = []
n_thread = self.n_thread
n_jobs = self.n_jobs
random_state = self.random_state
int_params = self.int_params
if scale_pos_use is True:
scale_pos_weight = round(np.sum(train_y == 0) / np.sum(train_y == 1), 1)
if 'scale_pos_weight' in self.params_dict[clf_name].keys():
self.params_dict[clf_name]['scale_pos_weight'] = \
tuple([x * scale_pos_weight for x in list(self.params_dict[clf_name]['scale_pos_weight'])])
else:
self.params_dict[clf_name]['scale_pos_weight'] = scale_pos_weight
else:
self.params_dict[clf_name]['scale_pos_weight'] = 1
print("[bayes_optimization] scale_pos_weight: {}".format(
self.params_dict[clf_name]['scale_pos_weight']), flush=True)
def _eval(**clf_args):
for param in int_params:
if clf_args.get(param):
clf_args[param] = int(round(clf_args[param]))
print("[eval] args: {}".format(clf_args))
if clf_name == 'lgb':
model = lgb.LGBMClassifier(**clf_args,
n_jobs=n_jobs,
random_state=random_state)
else:
model = xgb.XGBClassifier(**clf_args,
nthread=n_thread,
n_jobs=n_jobs,
random_state=random_state)
model.fit(train_x, train_y)
train_prob = model.predict_proba(train_x)[:, 1]
test_prob = model.predict_proba(test_x)[:, 1]
oot_prob = model.predict_proba(oot_x)[:, 1]
train_acc = round(accuracy_score(train_y, train_prob > 0.5), 5)
train_auc = round(roc_auc_score(train_y, train_prob), 5)
test_acc = round(accuracy_score(test_y, test_prob > 0.5), 5)
test_auc = round(roc_auc_score(test_y, test_prob), 5)
oot_acc = round(accuracy_score(oot_y, oot_prob > 0.5), 5)
oot_auc = round(roc_auc_score(oot_y, oot_prob), 5)
res_list.append({
'train_acc': train_acc, 'train_auc': train_auc,
'test_acc': test_acc, 'test_auc': test_auc,
'oot_acc': oot_acc, 'oot_auc': oot_auc
})
if target == 'auc':
target_value = test_auc
else:
target_value = test_acc
print('[train_acc] {}, [train_auc] {}, [test_acc] {}, [test_auc] {}, [oot_acc] {}, [oot_auc] {}'.format(
train_acc, train_auc, test_acc, test_auc, oot_acc, oot_auc), flush=True)
return target_value
print("[bayes_optimization] {}".format(self.params_dict[clf_name]), flush=True)
clf_bo = BayesianOptimization(f=_eval, pbounds=self.params_dict[clf_name])
clf_bo.maximize(**bayes_args)
self.best_param = clf_bo.max['params']
for param in int_params:
if self.params_dict[clf_name].get(param):
self.best_param[param] = int(round(self.best_param[param]))
self.best_param['nthread'] = n_thread
self.best_param['n_jobs'] = n_jobs
self.best_param['random_state'] = random_state
res_bo = []
for id, bo in enumerate(clf_bo.res):
for param in int_params:
if bo['params'].get(param):
bo['params'][param] = int(round(bo['params'][param]))
res_bo.append([bo.update(res_list[id])])
return clf_bo, self.best_param, res_bo
def eval(self, x, y):
# 使用模型对测试集数据进行预测
predictions = self.best_model.predict_proba(x)[:, 1]
print("[MAE] {}".format(mean_absolute_error(y_true=y, y_pred=predictions)), flush=True)
print("[ACC] {}".format(accuracy_score(y_true=y, y_pred=predictions > 0.5), flush=True))
print("[AUC] {}".format(roc_auc_score(y_true=y, y_score=predictions)), flush=True)
print("[confusion_matrix] \n{} ".format(confusion_matrix(y_true=y, y_pred=predictions > 0.5)), flush=True)
return predictions
def train(self, clf_name, x, y):
print("[train] best_param: {}".format(self.best_param), flush=True)
if clf_name == 'lgb':
print("[train] use model: LGBMClassifier", flush=True)
self.best_model = lgb.LGBMClassifier(**self.best_param)
else:
print("[train] use model: XGBClassifier", flush=True)
self.best_model = xgb.XGBClassifier(**self.best_param)
self.best_model.fit(x, y)
print("[train] best_model: {}".format(self.best_model), flush=True)
self.save_model()
模型训练示例
在进行模型训练的时候，就直接调用封装的GBDTModel类里面的函数即可，具体如下：
- 进行贝叶斯自动调参，并保存最优的参数。（数据均用pandas的DataFrame保存）
use_columns = ['col_1', 'col_2', 'col_3', 'col_4']
label_columns = ['label']
train_manager = GBDTModel(output_dir="./")
train_manager.bayes_optimization(
train_x=train_df[use_columns], train_y=np.array(train_df[label_columns]),
test_x=test_df[use_columns], test_y=test_df[label_columns],
oot_x=val_df[use_columns], oot_y=val_df[label_columns],
clf_name="xgb"
)
train_manager.save_best_param(file_name="params.json")
train_manager.read_best_param(file_name="params.json")
train_manager.train(clf_name="xgb",
x=train_data[use_columns], y=train_data[label_columns])
train_manager.eval(x=test_data[use_columns], y=test_data[label_columns])
train_manager.eval(x=val_data[use_columns], y=val_data[label_columns])
