1 from sklearn import tree,datasets 2 from sklearn.model_selection import train_test_split 3 wine=datasets.load_wine() 4 X,y=wine.data[:,:2],wine.target 5 X_train,X_test,y_train,y_test=train_test_split(X,y,random_state=1) 6 7 clf=tree.DecisionTreeClassifier(max_depth=3) 8 clf.fit(X_train,y_train) 9 print("the score of this model:{}".format(clf.score(X_test,y_test)))
1 from matplotlib.colors import ListedColormap 2 cmap_light=ListedColormap(['#FFAAAA','#AAFFAA','#AAAAFF']) 3 cmap_bold=ListedColormap(['#FF0000','#00FF00','#0000FF']) 4 5 import numpy as np 6 from matplotlib import pyplot as plt 7 x_min=X_train[:,0].min()-1 8 x_max=X_train[:,0].max()+1 9 y_min=X_train[:,1].min()-1 10 y_max=X_train[:,1].max()+1 11 xx,yy=np.meshgrid(np.arange(x_min,x_max,.02), 12 np.arange(y_min,y_max,.02)) 13 z=clf.predict(np.c_[xx.ravel(),yy.ravel()]) 14 z=z.reshape(xx.shape) 15 plt.figure() 16 plt.pcolormesh(xx,yy,z,cmap=cmap_light) 17 #plt.pcolormesh(xx,yy,z,cmap=plt.cm.Pastel1) 18 plt.scatter(X[:,0],X[:,1],c=y,cmap=cmap_bold,edgecolors='k',s=20) 19 plt.xlim(xx.min(),xx.max()) 20 plt.ylim(yy.min(),yy.max()) 21 plt.title("Decision Tree(max depth=3)") 22 plt.show()
1 import graphviz 2 from sklearn.tree import export_graphviz 3 export_graphviz(clf,out_file="wine.dot",class_names=wine.target_names, 4 feature_names=wine.feature_names[:2],impurity=False,filled=True) 5 with open("wine.dot") as f: 6 dot_graph = f.read() 7 dot=graphviz.Source(dot_graph) 8 dot.view()