forked from Kaggle/docker-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_lightgbm.py
More file actions
61 lines (51 loc) · 1.92 KB
/
Copy pathtest_lightgbm.py
File metadata and controls
61 lines (51 loc) · 1.92 KB
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
import unittest
import lightgbm as lgb
from common import gpu_test
class TestLightgbm(unittest.TestCase):
# Based on the "simple_example" from their documentation:
# https://github.com/Microsoft/LightGBM/blob/master/examples/python-guide/simple_example.py
def test_cpu(self):
lgb_train = lgb.Dataset('/input/tests/data/lgb_train.bin')
lgb_eval = lgb.Dataset('/input/tests/data/lgb_test.bin', reference=lgb_train)
params = {
'task': 'train',
'boosting_type': 'gbdt',
'objective': 'regression',
'metric': {'l2', 'auc'},
'num_leaves': 31,
'learning_rate': 0.05,
'feature_fraction': 0.9,
'bagging_fraction': 0.8,
'bagging_freq': 5,
'verbose': 0
}
# Run only one round for faster test
gbm = lgb.train(params,
lgb_train,
num_boost_round=1,
valid_sets=lgb_eval,
early_stopping_rounds=1)
self.assertEqual(1, gbm.best_iteration)
@gpu_test
def test_gpu(self):
lgb_train = lgb.Dataset('/input/tests/data/lgb_train.bin')
lgb_eval = lgb.Dataset('/input/tests/data/lgb_test.bin', reference=lgb_train)
params = {
'boosting_type': 'gbdt',
'objective': 'regression',
'metric': 'auc',
'num_leaves': 31,
'learning_rate': 0.05,
'feature_fraction': 0.9,
'bagging_fraction': 0.8,
'bagging_freq': 5,
'verbose': 1,
'device': 'gpu'
}
# Run only one round for faster test
gbm = lgb.train(params,
lgb_train,
num_boost_round=1,
valid_sets=lgb_eval,
early_stopping_rounds=1)
self.assertEqual(1, gbm.best_iteration)