Python Databricks环境下使用dtreeviz无法可视化决策树如何解决？

前置准备 #

首先在Databricks集群的库管理界面安装两个依赖包，安装完成后重启集群：

• dtreeviz（与你当前导入路径匹配的1.x版本）
• python-graphviz

问题原因 #

• 你当前代码创建DataFrame时重复声明了Feature02列，漏写了Feature03，导致传入模型的特征列不符合预期
• Databricks Notebook不支持直接渲染dtreeviz返回的DTreeViz实例，需要手动转换为SVG格式后调用HTML渲染接口展示

修复后完整代码 #

import pandas as pd
from sklearn import preprocessing, tree
from dtreeviz.trees import dtreeviz

Things = {'Feature01': [3,4,5,0], 
          'Feature02': [4,5,6,0], 
          'Feature03': [1,2,3,8], 
          'Target01': ['Red','Blue','Teal','Red']}
# 修正列名重复问题
df = pd.DataFrame(Things,
                  columns= ['Feature01', 'Feature02', 
                            'Feature03', 'Target01']) 

label_encoder = preprocessing.LabelEncoder()
label_encoder.fit(df.Target01)
df['target'] = label_encoder.transform(df.Target01)

classifier = tree.DecisionTreeClassifier()
classifier.fit(df.iloc[:,:3], df.target)

viz = dtreeviz(classifier,
         df.iloc[:,:3],
         df.target,
         target_name='toy',
         feature_names=df.columns[0:3],
         class_names=list(label_encoder.classes_)
         )
# 适配Databricks渲染逻辑
displayHTML(viz.svg())

内容的提问来源于stack exchange，提问作者Dario Federici