xgboost visualization with tensorboard

xgb_tb.py

'''
Updates:
1. Using tensorboardX to replace tensorboard_logger as I prefer tensorboardX's API
2. Log the tree growth process, which would be displayed under "IMAGES" tab in Tensorboard 
'''

import pandas as pd
import xgboost as xgb

from sklearn.model_selection import train_test_split
from sklearn.datasets import load_boston

from xgboost.callback import early_stop
from tensorboardX import SummaryWriter

import matplotlib.pyplot as plt

def save_ax(ax, filename, **kwargs):
    ax.axis("off")
    ax.figure.canvas.draw()
    trans = ax.figure.dpi_scale_trans.inverted() 
    bbox = ax.bbox.transformed(trans)
    plt.savefig(filename, dpi="figure", bbox_inches=bbox,  **kwargs)
    ax.axis("on")
    im = plt.imread(filename)
    return im

def logspy(comment):
    writer = SummaryWriter(comment=comment)
    
    def callback(env):
        writer.add_scalar('train', env.evaluation_result_list[0][1], env.iteration)
        writer.add_scalar('val', env.evaluation_result_list[1][1], env.iteration)
        writer.add_scalar('nodes', env.model.trees_to_dataframe().shape[0], env.iteration)
        arr = save_ax(xgb.plot_tree(env.model, num_trees=env.iteration), 'tmp.png')
        writer.add_images('tree', arr.reshape(-1, *arr.shape), env.iteration, dataformats='NHWC')
    return callback

boston = load_boston()
df = pd.DataFrame(boston.data, columns=boston.feature_names)

x1, x2, y1, y2 = train_test_split(
    df, boston.target, test_size=0.1, random_state=18)

dtrain = xgb.DMatrix(x1, y1)
dvalid = xgb.DMatrix(x2, y2)
watchlist = [(dtrain, 'train'), (dvalid, 'valid')]


model = xgb.train(
    params={},
    num_boost_round=100,
    dtrain=dtrain,
    evals=watchlist,
    callbacks=[logspy(''), early_stop(5)])