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()
