时间戳

首页 > 解决方案 > 在 Python 中将 create_tree_digraph 图写入 png 文件

python - 在 Python 中将 create_tree_digraph 图写入 png 文件

问题描述

我希望我的 lightgbm 模型树保存为 .png 格式。我尝试了 lightgbm API 的两种绘图方法 - plot_tree 和 create_tree_diagraph。

import lightgbm as lgb
from sklearn.datasets import load_iris

X, y = load_iris(True)
clf = lgb.LGBMClassifier()
clf.fit(X, y)

当我使用 plot_tree 时，它​​会显示树，但代替值的是小空白框

lgb.plot_tree(clf, tree_index=0)

当我尝试 create_tree_diagraph 时，我得到了图表，但我无法按原样保存它。

lgb.create_tree_digraph(clf)

我使用下面的代码将其保存为一个文件，但它被保存为第一个图（使用 plot_tree）

import graphviz
s = graphviz.Source(graph_b.source, filename = "test1.gv", format = "png")
s.view()

将绘图保存为图像的任何建议。我最终想把这些树图写成excel。我正在使用graphviz 0.8.3版谢谢，

标签： pythonplotlightgbm

解决方案

我想你正在使用 jupyter notebook 或类似的东西。这对我有用，但肯定不是最好的方法。

ax = lgb.create_tree_digraph(clf)
with open('fst.svg', 'w') as f:
    f.write(ax._repr_svg_())

推荐阅读