test_engine.py 10.1 KB
Newer Older
Guolin Ke's avatar
Guolin Ke committed
1
# coding: utf-8
wxchan's avatar
wxchan committed
2
# pylint: skip-file
wxchan's avatar
wxchan committed
3
4
5
6
7
import copy
import math
import os
import unittest

Guolin Ke's avatar
Guolin Ke committed
8
import lightgbm as lgb
wxchan's avatar
wxchan committed
9
10
11
12
import numpy as np
from sklearn.datasets import (load_boston, load_breast_cancer, load_digits,
                              load_iris)
from sklearn.metrics import log_loss, mean_absolute_error, mean_squared_error
wxchan's avatar
wxchan committed
13
from sklearn.model_selection import train_test_split, TimeSeriesSplit
wxchan's avatar
wxchan committed
14

wxchan's avatar
wxchan committed
15
16
17
18
19
20
try:
    import pandas as pd
    IS_PANDAS_INSTALLED = True
except ImportError:
    IS_PANDAS_INSTALLED = False

wxchan's avatar
wxchan committed
21
22
try:
    import cPickle as pickle
wxchan's avatar
wxchan committed
23
except ImportError:
wxchan's avatar
wxchan committed
24
    import pickle
wxchan's avatar
wxchan committed
25

wxchan's avatar
wxchan committed
26

wxchan's avatar
wxchan committed
27
28
29
def multi_logloss(y_true, y_pred):
    return np.mean([-math.log(y_pred[i][y]) for i, y in enumerate(y_true)])

wxchan's avatar
wxchan committed
30

wxchan's avatar
wxchan committed
31
32
33
34
class template(object):
    @staticmethod
    def test_template(params={'objective': 'regression', 'metric': 'l2'},
                      X_y=load_boston(True), feval=mean_squared_error,
Guolin Ke's avatar
Guolin Ke committed
35
                      num_round=200, init_model=None, custom_eval=None,
Guolin Ke's avatar
Guolin Ke committed
36
                      early_stopping_rounds=5,
wxchan's avatar
wxchan committed
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
                      return_data=False, return_model=False):
        params['verbose'], params['seed'] = -1, 42
        X_train, X_test, y_train, y_test = train_test_split(*X_y, test_size=0.1, random_state=42)
        lgb_train = lgb.Dataset(X_train, y_train, params=params)
        lgb_eval = lgb.Dataset(X_test, y_test, reference=lgb_train, params=params)
        if return_data:
            return lgb_train, lgb_eval
        evals_result = {}
        gbm = lgb.train(params, lgb_train,
                        num_boost_round=num_round,
                        valid_sets=lgb_eval,
                        valid_names='eval',
                        verbose_eval=False,
                        feval=custom_eval,
                        evals_result=evals_result,
                        early_stopping_rounds=early_stopping_rounds,
                        init_model=init_model)
        if return_model:
            return gbm
        else:
            return evals_result, feval(y_test, gbm.predict(X_test, gbm.best_iteration))
wxchan's avatar
wxchan committed
58

wxchan's avatar
wxchan committed
59

wxchan's avatar
wxchan committed
60
class TestEngine(unittest.TestCase):
wxchan's avatar
wxchan committed
61
62

    def test_binary(self):
wxchan's avatar
wxchan committed
63
        X_y = load_breast_cancer(True)
wxchan's avatar
wxchan committed
64
        params = {
wxchan's avatar
wxchan committed
65
66
            'objective': 'binary',
            'metric': 'binary_logloss'
wxchan's avatar
wxchan committed
67
        }
wxchan's avatar
wxchan committed
68
        evals_result, ret = template.test_template(params, X_y, log_loss)
wxchan's avatar
wxchan committed
69
        self.assertLess(ret, 0.15)
Guolin Ke's avatar
Guolin Ke committed
70
        self.assertAlmostEqual(min(evals_result['eval']['binary_logloss']), ret, places=5)
wxchan's avatar
wxchan committed
71

wxchan's avatar
wxchan committed
72
    def test_regreesion(self):
wxchan's avatar
wxchan committed
73
        evals_result, ret = template.test_template()
74
        self.assertLess(ret, 16)
wxchan's avatar
wxchan committed
75
76
77
78
79
        self.assertAlmostEqual(min(evals_result['eval']['l2']), ret, places=5)

    def test_multiclass(self):
        X_y = load_digits(10, True)
        params = {
wxchan's avatar
wxchan committed
80
81
82
            'objective': 'multiclass',
            'metric': 'multi_logloss',
            'num_class': 10
wxchan's avatar
wxchan committed
83
        }
wxchan's avatar
wxchan committed
84
        evals_result, ret = template.test_template(params, X_y, multi_logloss)
wxchan's avatar
wxchan committed
85
86
87
        self.assertLess(ret, 0.2)
        self.assertAlmostEqual(min(evals_result['eval']['multi_logloss']), ret, places=5)

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
    def test_early_stopping(self):
        X_y = load_breast_cancer(True)
        params = {
            'objective': 'binary',
            'metric': 'binary_logloss',
            'verbose': -1,
            'seed': 42
        }
        X_train, X_test, y_train, y_test = train_test_split(*X_y, test_size=0.1, random_state=42)
        lgb_train = lgb.Dataset(X_train, y_train)
        lgb_eval = lgb.Dataset(X_test, y_test, reference=lgb_train)
        # no early stopping
        gbm = lgb.train(params, lgb_train,
                        num_boost_round=10,
                        valid_sets=lgb_eval,
                        verbose_eval=False,
                        early_stopping_rounds=5)
        self.assertEqual(gbm.best_iteration, -1)
        # early stopping occurs
        gbm = lgb.train(params, lgb_train,
                        num_boost_round=100,
                        valid_sets=lgb_eval,
                        verbose_eval=False,
                        early_stopping_rounds=5)
        self.assertLessEqual(gbm.best_iteration, 100)

wxchan's avatar
wxchan committed
114
115
    def test_continue_train_and_other(self):
        params = {
wxchan's avatar
wxchan committed
116
117
            'objective': 'regression',
            'metric': 'l1'
wxchan's avatar
wxchan committed
118
119
        }
        model_name = 'model.txt'
wxchan's avatar
wxchan committed
120
        gbm = template.test_template(params, num_round=20, return_model=True, early_stopping_rounds=-1)
wxchan's avatar
wxchan committed
121
        gbm.save_model(model_name)
wxchan's avatar
wxchan committed
122
123
124
        evals_result, ret = template.test_template(params, feval=mean_absolute_error,
                                                   num_round=80, init_model=model_name,
                                                   custom_eval=(lambda p, d: ('mae', mean_absolute_error(p, d.get_label()), False)))
Guolin Ke's avatar
Guolin Ke committed
125
        self.assertLess(ret, 3.5)
wxchan's avatar
wxchan committed
126
127
128
129
130
131
132
133
134
135
        self.assertAlmostEqual(min(evals_result['eval']['l1']), ret, places=5)
        for l1, mae in zip(evals_result['eval']['l1'], evals_result['eval']['mae']):
            self.assertAlmostEqual(l1, mae, places=5)
        self.assertIn('tree_info', gbm.dump_model())
        self.assertIsInstance(gbm.feature_importance(), np.ndarray)
        os.remove(model_name)

    def test_continue_train_multiclass(self):
        X_y = load_iris(True)
        params = {
wxchan's avatar
wxchan committed
136
137
138
            'objective': 'multiclass',
            'metric': 'multi_logloss',
            'num_class': 3
wxchan's avatar
wxchan committed
139
        }
wxchan's avatar
wxchan committed
140
141
142
        gbm = template.test_template(params, X_y, num_round=20, return_model=True, early_stopping_rounds=-1)
        evals_result, ret = template.test_template(params, X_y, feval=multi_logloss,
                                                   num_round=80, init_model=gbm)
wxchan's avatar
wxchan committed
143
144
145
146
        self.assertLess(ret, 1.5)
        self.assertAlmostEqual(min(evals_result['eval']['multi_logloss']), ret, places=5)

    def test_cv(self):
wxchan's avatar
wxchan committed
147
        lgb_train, _ = template.test_template(return_data=True)
wxchan's avatar
wxchan committed
148
        lgb.cv({'verbose': -1}, lgb_train, num_boost_round=20, nfold=5, shuffle=False,
149
               metrics='l1', verbose_eval=False,
150
151
152
               callbacks=[lgb.reset_parameter(learning_rate=lambda i: 0.1 - 0.001 * i)])
        lgb.cv({'verbose': -1}, lgb_train, num_boost_round=20, nfold=5, shuffle=True,
               metrics='l1', verbose_eval=False,
153
               callbacks=[lgb.reset_parameter(learning_rate=lambda i: 0.1 - 0.001 * i)])
wxchan's avatar
wxchan committed
154
155
156
        tss = TimeSeriesSplit(3)
        lgb.cv({'verbose': -1}, lgb_train, num_boost_round=20, data_splitter=tss, nfold=5,  # test if wrong nfold is ignored
               metrics='l2', verbose_eval=False)
wxchan's avatar
wxchan committed
157

wxchan's avatar
wxchan committed
158
159
160
161
162
163
    def test_feature_name(self):
        lgb_train, _ = template.test_template(return_data=True)
        feature_names = ['f' + str(i) for i in range(13)]
        gbm = lgb.train({'verbose': -1}, lgb_train, num_boost_round=10, feature_name=feature_names)
        self.assertListEqual(feature_names, gbm.feature_name())

wxchan's avatar
wxchan committed
164
    def test_save_load_copy_pickle(self):
wxchan's avatar
wxchan committed
165
166
        gbm = template.test_template(num_round=20, return_model=True)
        _, ret_origin = template.test_template(init_model=gbm)
wxchan's avatar
wxchan committed
167
168
        other_ret = []
        gbm.save_model('lgb.model')
wxchan's avatar
wxchan committed
169
        other_ret.append(template.test_template(init_model='lgb.model')[1])
wxchan's avatar
wxchan committed
170
        gbm_load = lgb.Booster(model_file='lgb.model')
wxchan's avatar
wxchan committed
171
172
173
        other_ret.append(template.test_template(init_model=gbm_load)[1])
        other_ret.append(template.test_template(init_model=copy.copy(gbm))[1])
        other_ret.append(template.test_template(init_model=copy.deepcopy(gbm))[1])
wxchan's avatar
wxchan committed
174
175
176
177
        with open('lgb.pkl', 'wb') as f:
            pickle.dump(gbm, f)
        with open('lgb.pkl', 'rb') as f:
            gbm_pickle = pickle.load(f)
wxchan's avatar
wxchan committed
178
        other_ret.append(template.test_template(init_model=gbm_pickle)[1])
wxchan's avatar
wxchan committed
179
        gbm_pickles = pickle.loads(pickle.dumps(gbm))
wxchan's avatar
wxchan committed
180
        other_ret.append(template.test_template(init_model=gbm_pickles)[1])
wxchan's avatar
wxchan committed
181
182
        for ret in other_ret:
            self.assertAlmostEqual(ret_origin, ret, places=5)
wxchan's avatar
wxchan committed
183

184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
    @unittest.skipIf(not IS_PANDAS_INSTALLED, 'pandas not installed')
    def test_pandas_categorical(self):
        X = pd.DataFrame({"A": np.random.permutation(['a', 'b', 'c', 'd'] * 75),  # str
                          "B": np.random.permutation([1, 2, 3] * 100),  # int
                          "C": np.random.permutation([0.1, 0.2, -0.1, -0.1, 0.2] * 60),  # float
                          "D": np.random.permutation([True, False] * 150)})  # bool
        y = np.random.permutation([0, 1] * 150)
        X_test = pd.DataFrame({"A": np.random.permutation(['a', 'b', 'e'] * 20),
                               "B": np.random.permutation([1, 3] * 30),
                               "C": np.random.permutation([0.1, -0.1, 0.2, 0.2] * 15),
                               "D": np.random.permutation([True, False] * 30)})
        for col in ["A", "B", "C", "D"]:
            X[col] = X[col].astype('category')
            X_test[col] = X_test[col].astype('category')
        params = {
            'objective': 'binary',
            'metric': 'binary_logloss',
            'verbose': -1
        }
        lgb_train = lgb.Dataset(X, y)
        gbm0 = lgb.train(params, lgb_train, num_boost_round=10, verbose_eval=False)
        pred0 = list(gbm0.predict(X_test))
        lgb_train = lgb.Dataset(X, y)
        gbm1 = lgb.train(params, lgb_train, num_boost_round=10, verbose_eval=False,
                         categorical_feature=[0])
        pred1 = list(gbm1.predict(X_test))
        lgb_train = lgb.Dataset(X, y)
        gbm2 = lgb.train(params, lgb_train, num_boost_round=10, verbose_eval=False,
                         categorical_feature=['A'])
        pred2 = list(gbm2.predict(X_test))
        lgb_train = lgb.Dataset(X, y)
        gbm3 = lgb.train(params, lgb_train, num_boost_round=10, verbose_eval=False,
                         categorical_feature=['A', 'B', 'C', 'D'])
        pred3 = list(gbm3.predict(X_test))
        lgb_train = lgb.Dataset(X, y)
        gbm3.save_model('categorical.model')
        gbm4 = lgb.Booster(model_file='categorical.model')
        pred4 = list(gbm4.predict(X_test))
222
223
224
225
        np.testing.assert_almost_equal(pred0, pred1)
        np.testing.assert_almost_equal(pred0, pred2)
        np.testing.assert_almost_equal(pred0, pred3)
        np.testing.assert_almost_equal(pred0, pred4)
226

wxchan's avatar
wxchan committed
227

wxchan's avatar
wxchan committed
228
229
230
print("----------------------------------------------------------------------")
print("running test_engine.py")
unittest.main()