我相信这个答案比这里的其他答案更正确:
from sklearn.tree import _tree
def tree_to_code(tree, feature_names):
tree_ = tree.tree_
feature_name = [
feature_names[i] if i != _tree.TREE_UNDEFINED else "undefined!"
for i in tree_.feature
]
print "def tree({}):".format(", ".join(feature_names))
def recurse(node, depth):
indent = " " * depth
if tree_.feature[node] != _tree.TREE_UNDEFINED:
name = feature_name[node]
threshold = tree_.threshold[node]
print "{}if {} <= {}:".format(indent, name, threshold)
recurse(tree_.children_left[node], depth + 1)
print "{}else: # if {} > {}".format(indent, name, threshold)
recurse(tree_.children_right[node], depth + 1)
else:
print "{}return {}".format(indent, tree_.value[node])
recurse(0, 1)
这将打印一个有效的Python函数。这里是一个树的示例输出,试图返回其输入,数字在0和10之间。
def tree(f0):
if f0 <= 6.0:
if f0 <= 1.5:
return [[ 0.]]
else: # if f0 > 1.5
if f0 <= 4.5:
if f0 <= 3.5:
return [[ 3.]]
else: # if f0 > 3.5
return [[ 4.]]
else: # if f0 > 4.5
return [[ 5.]]
else: # if f0 > 6.0
if f0 <= 8.5:
if f0 <= 7.5:
return [[ 7.]]
else: # if f0 > 7.5
return [[ 8.]]
else: # if f0 > 8.5
return [[ 9.]]
这里有一些绊脚石我在其他答案看到:
>使用tree_.threshold == -2来决定节点是否是一个叶子不是一个好主意。如果它是一个阈值为-2的真正的决策节点呢?相反,你应该看看tree.feature或tree.children_ *。> line feature = [feature_names [i] for i in tree_.feature]与我的sklearn版本崩溃,因为tree.tree_feature的一些值是-2(特别是对于叶节点)。>没有必要在递归函数中有多个if语句,只有一个很好。