Visualizing decision tree in scikit-learn

后端 未结 11 2196
广开言路
广开言路 2020-12-02 10:26

I am trying to design a simple Decision Tree using scikit-learn in Python (I am using Anaconda\'s Ipython Notebook with Python 2.7.3 on Windows OS) and visualize it as follo

11条回答
  •  甜味超标
    2020-12-02 11:20

    The following also works fine:

    from sklearn.datasets import load_iris
    iris = load_iris()
    
    # Model (can also use single decision tree)
    from sklearn.ensemble import RandomForestClassifier
    model = RandomForestClassifier(n_estimators=10)
    
    # Train
    model.fit(iris.data, iris.target)
    # Extract single tree
    estimator = model.estimators_[5]
    
    from sklearn.tree import export_graphviz
    # Export as dot file
    export_graphviz(estimator, out_file='tree.dot', 
                    feature_names = iris.feature_names,
                    class_names = iris.target_names,
                    rounded = True, proportion = False, 
                    precision = 2, filled = True)
    
    # Convert to png using system command (requires Graphviz)
    from subprocess import call
    call(['dot', '-Tpng', 'tree.dot', '-o', 'tree.png', '-Gdpi=600'])
    
    # Display in jupyter notebook
    from IPython.display import Image
    Image(filename = 'tree.png')
    

    You can find the source here

提交回复
热议问题