forked from mlflow/mlflow
-
Notifications
You must be signed in to change notification settings - Fork 0
/
utils.py
110 lines (89 loc) · 3.15 KB
/
utils.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
import importlib
import json
import os.path
import warnings
import numpy as np
def _label_encoder_to_json(le):
"""Returns a JSON compatible dictionary"""
meta = {}
for k, v in le.__dict__.items():
if isinstance(v, np.ndarray):
meta[k] = v.tolist()
else:
meta[k] = v
return meta
def _label_encoder_from_json(doc):
"""Load the encoder back from a JSON compatible dict"""
from lightgbm.compat import _LGBMLabelEncoder
le = _LGBMLabelEncoder()
meta = {}
for k, v in doc.items():
if k == "classes_":
le.classes_ = np.array(v) if v is not None else None
continue
meta[k] = v
le.__dict__.update(meta)
return le
def _save_lgb_attr(model_dir, fname, attr_dict):
with open(os.path.join(model_dir, "{}.json".format(fname)), "w") as f:
json.dump(attr_dict, f)
def _load_lgb_attr(model_dir, fname):
try:
with open(os.path.join(model_dir, "{}.json".format(fname))) as f:
attr = json.load(f)
return attr
except IOError:
return None
def _save_lgb_model(lgb_model, model_path) -> None:
import lightgbm as lgb
model_dir = os.path.dirname(model_path)
if not isinstance(lgb_model, lgb.Booster):
meta = {}
for k, v in lgb_model.__dict__.items():
if k == "_le":
meta["_le"] = _label_encoder_to_json(v) if v else None
continue
if k == "_Booster":
continue
if k == "_classes" and v is not None:
meta["_classes"] = v.tolist()
continue
if k == "_class_map" and v:
py_dict = {}
for clazz, encoded in v.items():
py_dict[int(clazz)] = int(encoded)
v = py_dict
try:
json.dumps({k: v})
meta[k] = v
except TypeError:
warnings.warn(str(k) + " is not saved in Scikit-Learn meta.", UserWarning)
_save_lgb_attr(model_dir, "scikit-learn", meta)
lgb_model = lgb_model._Booster
lgb_model.save_model(model_path)
_save_lgb_attr(model_dir, "params", lgb_model.params)
def _load_lgb_model(lgb_model_class, model_path):
import lightgbm as lgb
module, cls = lgb_model_class.rsplit(".", maxsplit=1)
model_dir = os.path.dirname(model_path)
sk_attr = _load_lgb_attr(model_dir, "scikit-learn")
bst_params = _load_lgb_attr(model_dir, "params")
booster = lgb.Booster(model_file=model_path, params=bst_params)
if sk_attr is None:
warnings.warn("Loading a native LightGBM model with Scikit-Learn interface.")
return booster
sk_model = getattr(importlib.import_module(module), cls)()
states = {}
for k, v in sk_attr.items():
if k == "_le":
sk_model._le = _label_encoder_from_json(v)
continue
if k == "_classes":
sk_model._classes = np.array(v)
continue
states[k] = v
sk_model.__dict__.update(states)
# Delete the attribute after load
booster.set_attr(scikit_learn=None)
sk_model._Booster = booster
return sk_model