赞
踩
from sklearn.datasets import load_iris from sklearn.tree import DecisionTreeClassifier from sklearn.model_selection import train_test_split from sklearn.tree import export_graphviz from IPython.display import Image import matplotlib.pyplot as plt import pydotplus # 1.定义X和y X = df_churn.iloc[:, :-1] y = df_churn.Status # 2. 建立决策树模型 dt_model = DecisionTreeClassifier(max_depth=30, min_samples_split=50, min_samples_leaf=25, max_leaf_nodes=100, class_weight='balanced', ccp_alpha=0.0001) # Fit model to training data dt_model.fit(X_train, y_train) # 3. 决策树的可视化 tmp_dot_file = 'decision_tree_tmp.dot' export_graphviz(dt_model, out_file=tmp_dot_file,filled=True,feature_names=X.columns, class_names=list(set(y)),impurity=False) with open(tmp_dot_file) as f: dot_graph = f.read() graph = pydotplus.graph_from_dot_data(dot_graph) graph.write_pdf('example.pdf') #保存图像为pdf格式 Image(graph.create_png()) #绘制图像为png格式
就可以得到这样的结果啦:
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。