Skip to content

GitLab

Commit 7f94fd9c authored by wxchan's avatar wxchan Committed by Guolin Ke
Browse files

fix feature importance (#468)

parent edd618bf
Hide whitespace changes
Inline Side-by-side
Showing with 5 additions and 4 deletions
python-package/lightgbm/basic.py
View file @ 7f94fd9c
...@@ -1653,10 +1653,11 @@ class Booster(object): ...@@ -1653,10 +1653,11 @@ class Booster(object):
def dfs(root): def dfs(root):
if "split_feature" in root: if "split_feature" in root:
if importance_type == 'split': if root['split_gain'] > 0:
ret[root["split_feature"]] += 1 if importance_type == 'split':
elif importance_type == 'gain': ret[root["split_feature"]] += 1
ret[root["split_feature"]] += root["split_gain"] elif importance_type == 'gain':
ret[root["split_feature"]] += root["split_gain"]
dfs(root["left_child"]) dfs(root["left_child"])
dfs(root["right_child"]) dfs(root["right_child"])
for tree in dump_model["tree_info"]: for tree in dump_model["tree_info"]:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment