Writing create_tree_digraph plot to a png file in Python

Question

I want the tree of my lightgbm model to save to a .png format. I have tried two plotting methods from lightgbm API - plot_tree and 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)

When I use plot_tree, it displays the tree but in place of values there are small blank boxes

lgb.plot_tree(clf, tree_index=0)

When I try the create_tree_diagraph, I get the graph but I cant save it as it is.

lgb.create_tree_digraph(clf)

I used the below code to save it a file but that gets saved as the first plot (using plot_tree)

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

Any suggestions to save the plot as an image. I ultimately want to write these tree plots to excel. I am using graphviz version 0.8.3 Thanks,

Answer 1

I suppose you are using a jupyter notebook or something like that. This worked for me, but surely is not the best way.

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