GBDT算法实践-XGBoost模型可视化

上篇讲解了GBDT算法的实现，我们需要对模型结果进行可视化。注意基于Spark版本的模型存储需要调用model.nativeBooster.saveModel，保证模型文件通用性。分布式文件无法在线上预测中使用，也不方便模型可视化。详情参考我的GBDT算法原理分析，模型可视化展示决策树。

考虑到实际生产中需要显示特征名称，而不仅仅是特征编号，源码中新增了featureconf配置，可以更好地展示决策树。

import xgboost as xgb
import codecs
import os
import sys
import getopt


def main(argv):
    modelfile = ""
    featureconf = ""

    try:
        opts, args = getopt.getopt(argv, "hm:f:", ["modelfile=", "featureconf="])
    except getopt.GetoptError:
        print
        'plotTree.py -m <modelfile> -f <featureconf> '
    for opt, arg in opts:
        if opt == '-h':
            print
            'plotTree.py -m <modelfile> -f <featureconf>'
            sys.exit()
        elif opt in ("-m", "--modelfile"):
            modelfile = arg
        elif opt in ("-f", "--featureconf"):
            featureconf = arg

    fmap = {}
    if os.path.exists(modelfile):
        if os.path.exists(featureconf):
            if os.path.exists("./featuremap.txt"):
                os.remove("./featuremap.txt")
            filemap = open("./featuremap.txt", "wb")
            filemap.write(("0 null q" + '\n').encode())

            for line in open(featureconf, "rb"):
                # print(line)
                content = line.decode().strip('\n').split('\t')
                # print(content)
                if len(content) >= 3:
                    fmap[content[1]] = content[2]
                    temp = (content[1] + " " + content[2] + " " + "q" + '\n').encode()
                    filemap.write(temp)
            filemap.close()
        xgb_model = xgb.Booster()
        xgb_model.load_model(modelfile)
        score = xgb_model.get_fscore()
        importance = {}
        for key, value in score.items():
            temp = key[1:None]
            if temp in fmap:
                importance[fmap[temp]] = int(value)
            else:
                importance[key] = int(value)
        sort_importance = sorted(importance.items(), key=lambda importance: importance[1])

        for values in sort_importance[::-1]:
            print(values)
        tree = xgb_model.get_dump()
        i = 0
        while i < len(tree):
            f = codecs.open(modelfile + "_" + str(i)+'.png', mode='wb')
            g = xgb.to_graphviz(xgb_model, './featuremap.txt', i)
            f.write(g.pipe('png'))
            f.close()
            i = i+1
        print("success!")


if __name__ == "__main__":
    main(sys.argv[1:])

实际效果如下：