dtree_sklearn.py 文件源码

def visualize_tree(tree, feature_name, dot_file):
    """Create tree png using graphviz.
    tree -- scikit-learn DecsisionTree.
    feature_names -- list of feature names.
    dot_file -- dot file name and path
    """
    with open("tree.dot", 'w') as f:
        export_graphviz(tree, out_file=f,
                        feature_names=feature_name)

    dt_png = dot_file.replace('dot', 'png')
    command = ["dot", "-Tpng", dot_file, "-o", dt_png]
    try:
        subprocess.check_call(command)
    except Exception as e:
        print e
        exit("Could not run dot, ie graphviz, to "
             "produce visualization")