plot_example.py 1.06 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
# coding: utf-8
# pylint: disable = invalid-name, C0111
import lightgbm as lgb
import pandas as pd

try:
    import matplotlib.pyplot as plt
except ImportError:
    raise ImportError('You need to install matplotlib for plot_example.py.')

# load or create your dataset
print('Load data...')
df_train = pd.read_csv('../regression/regression.train', header=None, sep='\t')

y_train = df_train[0]
X_train = df_train.drop(0, axis=1)

# create dataset for lightgbm
lgb_train = lgb.Dataset(X_train, y_train)

# specify your configurations as a dict
params = {
wxchan's avatar
wxchan committed
23
    'num_leaves': 5,
24
25
26
27
28
29
30
    'verbose': 0
}

print('Start training...')
# train
gbm = lgb.train(params,
                lgb_train,
wxchan's avatar
wxchan committed
31
32
33
                num_boost_round=100,
                feature_name=['f' + str(i + 1) for i in range(28)],
                categorical_feature=[21])
34
35
36
37
38

print('Plot feature importances...')
# plot feature importances
ax = lgb.plot_importance(gbm, max_num_features=10)
plt.show()
wxchan's avatar
wxchan committed
39
40
41
42
43

print('Plot 84th tree...')
# plot tree
lgb.plot_tree(gbm, tree_index=83, figsize=(20, 8), show_info=['split_gain'])
plt.show()