python绘制决策树

python 复制代码
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格式

就可以得到这样的结果啦:

相关推荐
EveryPossible几秒前
静态箭头连线
开发语言·javascript·ecmascript
NiKo_W3 分钟前
Git 版本回退与撤销修改
开发语言·git·安全
listhi52010 分钟前
Map对象在JavaScript循环中的使用
开发语言·前端·javascript
大可门耳23 分钟前
Qt第一课:Qt是什么?相对于其他框架的优劣势是什么
开发语言·qt
西阳未落1 小时前
C语言中的内存函数(memcpy, memmove, memcmp, memset)
c语言·开发语言
xchenhao3 小时前
SciKit-Learn 全面分析分类任务 breast_cancer 数据集
python·机器学习·分类·数据集·scikit-learn·svm
axban3 小时前
QT M/V架构开发实战:QFileSystemModel介绍
开发语言·qt·架构
钢门狂鸭5 小时前
关于rust的crates.io
开发语言·后端·rust
Lionel_SSL6 小时前
《深入理解Java虚拟机》第三章读书笔记:垃圾回收机制与内存管理
java·开发语言·jvm
技术猿188702783516 小时前
PHP 与 WebAssembly 的 “天然隔阂”
开发语言·php·wasm