import os
import time
import pydotplus
import numpy as np
import pandas as pd
from sklearn import tree
from sklearn.externals.six import StringIO
from sklearn.model_selection import train_test_split
print('Step 1.Loading data...')
data = pd.read_csv("/users/dzh/desktop/newtrain2.csv")
x = data.iloc[:,1:4].values
y = data.iloc[:,0].values
X_train,X_test,Y_train,Y_test = train_test_split(x,y,test_size=0.4)
print('---Loading and splitting completed.')
print('Step 2.Training...')
startTime = time.time()
clf = tree.DecisionTreeClassifier()
clf = clf.fit(X_train,Y_train)
print('---Training Completed.Took %f s.'%(time.time()-startTime))
feature_name = ['IS_LOCAL','HAS_FUND','type_dw']
target_name = ['0','1']
dot_data = StringIO()
tree.export_graphviz(clf,out_file = dot_data,feature_names=feature_name,
class_names=target_name,filled=True,rounded=True,
special_characters=True)
graph = pydotplus.graph_from_dot_data(dot_data.getvalue())
graph.write_pdf("/users/dzh/desktop/88Tree.pdf")
print('Visible tree plot saved as pdf.')
Python:CART算法可视化决策树
最新推荐文章于 2024-05-16 21:14:39 发布