-
Notifications
You must be signed in to change notification settings - Fork 0
/
models.py
115 lines (90 loc) · 3.25 KB
/
models.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
import datetime
import mlflow
import lightgbm
import xgboost as xgb
import os
def convert_str_to_num(var:str):
if var.isnumeric():
return int(var)
try:
return float(var)
except:
return var
class Model(object):
def __init__(self,args):
pass
def save(self,directory,use_mlflow):
raise NotImplementedError
def load(self,directory):
raise NotImplementedError
def predict(self,X):
raise NotImplementedError
def fit(self,X,y):
raise NotImplementedError
class LGBModel(object):
def __init__(self,args={}):
super().__init__()
self.model_name = 'model.lgb'
default_params = {
'num_leaves':31,
'max_depth':-1,
'learning_rate':0.1,
'n_estimators':1000,
'objective':'binary',
'reg_alpha':0.,
'reg_lambda':0.,
'subsample':1.,
'device':'gpu',
'random_state':datetime.datetime.now().microsecond,
}
self.params = {**default_params,**args}
self.model = lightgbm.LGBMRegressor(**self.params)
def save(self,directory,use_mlflow=True):
fullpath = os.path.join(directory,self.model_name)
if use_mlflow:
mlflow.lightgbm.save_model(self.model.booster_, fullpath)
else:
self.model.booster_.save_model(fullpath)
def load(self,directory,params=None):
fullpath = os.path.join(directory, self.model_name)
self.model = lightgbm.Booster(model_file=fullpath)
if params:
self.params = {k:convert_str_to_num(v) for k,v in params.items()}
def predict(self,X):
return self.model.predict(X)
def fit(self,X,y):
self.params['random_state'] = datetime.datetime.now().microsecond
self.model = lightgbm.LGBMRegressor(**self.params)
self.model.fit(X,y,eval_metric=['auc'])
class XGBModel(object):
def __init__(self,args={}):
super().__init__()
self.model_name = 'model.xgb'
default_params = {
'n_estimators':1000,
'max_depth':6,
'tree_method':'gpu_hist',
'use_label_encoder':False,
'subsample':0.5,
'objective':'binary:logistic',
'seed':datetime.datetime.now().microsecond,
}
self.params = {**default_params,**args}
self.model = xgb.XGBRegressor(**self.params)
def save(self,directory,use_mlflow=True):
fullpath = os.path.join(directory,self.model_name)
if use_mlflow:
mlflow.xgboost.save_model(self.model, fullpath)
else:
self.model.save_model(fullpath)
def load(self,directory,params=None):
fullpath = os.path.join(directory, self.model_name)
self.model.load_model(fullpath)
if params:
self.params = {k:convert_str_to_num(v) for k,v in params.items()}
def predict(self,X):
return self.model.predict(X)
def fit(self,X,y):
self.params['seed'] = datetime.datetime.now().microsecond
self.model = xgb.XGBRegressor(**self.params)
self.model.fit(X,y,eval_metric=['auc'])