使用dtreeviz可视化XGBClassifier模型时遇AttributeError求助

问题重现 #

训练并加载XGBClassifier模型后，使用dtreeviz绘制可视化图时触发错误：

AttributeError: 'Index' object has no attribute 'index'

原代码：

from dtreeviz.trees import *
viz = dtreeviz(loaded_model, X_train, y_train, tree_index = 10,feature_names = X_train.columns,
      class_names = ['Absence', 'Presence'], target_name = 'Heart Disease')
viz.view()

完整报错栈：

---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-18-2a3024516ed1> in <module>
  1 from dtreeviz.trees import *
  ----> 2 viz = dtreeviz(loaded_model, X_train, y_train, tree_index = 10, feature_names = X_train.columns,
  3               class_names = ['Absence', 'Presence'], target_name = 'Heart Disease')
  4 viz.view()

~\anaconda3\lib\site-packages\dtreeviz\trees.py in dtreeviz(tree_model, x_data, y_data, feature_names, target_name, class_names, tree_index, precision, orientation, instance_orientation, show_root_edge_labels, show_node_labels, show_just_path, fancy, histtype, highlight_path, X, max_X_features_LR, max_X_features_TD, depth_range_to_display, label_fontsize, ticks_fontsize, fontname, title, title_fontsize, colors, cmap, scale)
816     if shadow_tree.is_classifier():
817         nbins = get_num_bins(histtype, n_classes)
--> 818         node_heights = shadow_tree.get_split_node_heights(X_data, y_data, nbins=nbins)
819 
820     internal = []

~\anaconda3\lib\site-packages\dtreeviz\models\shadow_decision_tree.py in get_split_node_heights(self, X_train, y_train, nbins)
273         for node in self.internal:
274             # print(node.feature_name(), node.id)
--> 275             X_feature = X_train[:, node.feature()]
276             overall_feature_range = (np.min(X_feature), np.max(X_feature))
277             # print(f"range {overall_feature_range}")

~\anaconda3\lib\site-packages\dtreeviz\models\shadow_decision_tree.py in feature(self)
506         """Returns feature index used at this node"""
507 
--> 508         return self.shadow_tree.get_node_feature(self.id)
509 
510     def feature_name(self) -> (str, None):

~\anaconda3\lib\site-packages\dtreeviz\models\xgb_decision_tree.py in get_node_feature(self, id)
 76         feature_name = self._get_nodes_values("Feature")[id]
 77         try:
--> 78             return self.feature_names.index(feature_name)
 79         except ValueError as error:
 80             return self.__class__.NO_FEATURE

 AttributeError: 'Index' object has no attribute 'index'

问题原因 #

报错核心在于：X_train.columns是pandas Index对象，而dtreeviz处理XGBoost模型的逻辑中，调用了仅属于Python内置列表的index()方法，Index对象不支持该方法，因此触发AttributeError。

解决方案 #

将feature_names参数从pandas Index转换为普通Python列表，两种可行方式：

1. 使用tolist()方法转换
2. 使用list()强制转换

修改后的代码：

from dtreeviz.trees import *
viz = dtreeviz(loaded_model, X_train, y_train, tree_index=10,
               feature_names=X_train.columns.tolist(),  # 转换为列表
               class_names=['Absence', 'Presence'], target_name='Heart Disease')
viz.view()

验证逻辑 #

转换后，feature_names变为标准Python列表，调用index()方法可正常查找特征名称对应的索引，从而避免错误。

内容的提问来源于stack exchange，提问作者Nilotpole Kalita