基于XGB单机训练VS基于SPARK并行预测（XGBoost4j-spark无痛人流解决方案）

作者·黄崇远

『数据虫巢』

全文共9000字

题图ssyer.com

“ Python训练VS并行化预测，无痛人流般的解决方案，你值得拥有。”

理解本文需要有一定的技术基础，包括对于Xgboost的基本理解以及使用经验，基本的Spark开发能力，如果对于Xgboost4j-spark有一定的了解就更好了。

01

诉求背景

在这里，我就不做XGB的科普了，如果不清楚的，请自行谷歌，官网。

首先我们确定我们需要做的事情，那就是尝试在Python单机的环境下训练模型，获取到模型文件，然后加载在Spark环境中做并行预测，涉及到并行预测会用到XGBoost4j-spark框架。

这看起来是一个伪诉求，为什么会存在使用单机来训练，然后跑到Spark上预测的这种诉求，比如存在以下几个问题。

(1) 为什么不直接在单机，比如Python的XGB上进行Train以及Predict？

(2) 为什么不直接在XGBoost4j-spark上做Train以及Predict？

上面两个问题简直就是灵魂拷问，看似合情合理，无法推翻。来，让我们来逐一探讨一下。

• 先说单机Python的XGB。

我们使用XGB通常是分类或者回归场景，是一种相对偏传统的做法（非深度学习系列），所以带标的数据样本量级通常不会太大，几万是常见的，几十万也能接受，上百万也能处理的了（单机的内存和核稍微大点），所以在Python单机上训练是没有太多压力的。

并且，Python本来就是一个脚本式的语言，所以代码非常简洁，跑起来非常快，离线部署起来也不难。由于Python对应的很多数据处理的相关库，对于数据探索，特征挖掘，进而进行模型调优是非常便捷的。

如果是处于一种正常迭代的情况下，你可能会处于反复调整采样方式，不断的增减模型，不断的调参的这种循环中，这需要一个轻量级并且灵活的环境来支持，而Python的环境恰巧非常符合。这意味着Python对应的XGB环境是比较利于这种良好迭代节奏的。

说完了训练，说预测，并且这里说的不是服务化实时预测，说的是批量离线预测的话题，因为如果是实时，就不存在使不使用Spark做并行预测的话题了。

如果说我们的预测场景是几十上百万，又甚至是上千万的量级，其实单机都能撸的过来，再不行就做多线程嘛（Python无法做多线程，但是逻辑里可以做多进程的方式来实现并行），再不行就切割数据，分散到多个节点嘛，最后再聚拢数据。

我曾把模型利用自己写的调度脚本进行并行化，并将预处理和预测错峰执行，再同时并行三个模型预测，调调一些调度参数，把128GB内存，以及64核的单机服务器打的满满的，2个小时3个模型分别做1亿多数据的二分类。

OK，这里是1亿数据，如果是2亿，3亿，5亿呢？然后如果不止3个模型，是十个八个模型呢？解决肯定是可以解决的，但是做资源分配，并行处理，甚至多机拆解任务会把人搞死。关键一旦跑起来之后这个机器就基本上干不了其他的了，这意味着这压根儿不存在啥资源调度的问题。

所以，需要解决这种大规模多模型预测对于资源的消耗问题，甚至是效率问题。

• 说完了单机，说Spark的多机，指的是XGB的spark框架代表XGBoost4j-spark。

模型并行化，好处当然大家都知道的，大规模的数据不管是训练还是预测，都“咻咻”的，并且不需要考虑训练或者预测资源的问题，资源不够多配点excute就好了，别说3亿数据，只要集群中有资源，10亿我都给你分分钟预测出来。

所以预测这层天然是符合这种大规模离线预测场景的。

回到训练，除了上述单机场景中的迭代领域这个优点也就是对应Spark的缺点。spark其实相对来说是一个比较笨重的框架，任务提交、任务的响应和执行需要比较费时，如果资源临时不够还得排排队，这对于我们需要快速灵活迭代模型的诉求是相悖的。

除了上面的缺点还有缺点吗？有的，XGBoost4j-spark的训练过程在数据量少的情况下，其训练带来的精度有可能是低于单机的，相当于进一步稀释了训练样本，这在于训练样本数本来就不算特别多的场景中，是有一定的影响的。

关于这一点，我们在后面拆解XGBoost4j-spark源码的时候再来进一步说明。

所以，在这里，在训练数据量百万级以内，离线预测量级在数亿以上的场景中，单机训练Spark并行预测的搭配太合适了，简直就是天造地设的...

好吧，差点出口成章了。既然如此般配，那么可以直接Python版的XGB训练好的Model，直接丢到XGBoost4j-spark中load，然后愉快的预测呢？

答案是不行！竟然是不行。我也很纳闷，为啥不行。做个工作流每天让Python定时训练模型，然后丢到Spark环境中每天做预测，然后还一边用Python来调模型，一旦模型指标非常OK的，离线验证之后，直接丢到Spark流程中替换。

这个过程是多么的自然，但结果竟然是不行。说好的XGBoost4j-spark是Xgboost的分支项目的呢，这也不像亲儿子啊。

真的是不行，所以才有了这个文章，和研究方向。我试图从源码中探索为什么不行，理论上一定行的事，然后找到这个路径的解决方案。

02

XGBoost单机源码拆解

代码来源：git clone --recursivehttps://github.com/dmlc/xgboost

这是XGB在github上的开源代码，其中也包括了XGBoost4j-spark这个分支的项目代码，一箭双雕。

首先来看下XGB主体部分的代码目录：

|--xgboost

    |--include

        |--xgboost //定义了 xgboost 相关的头文件

    |--src

        |--c_api              

        |--common //一些通用文件，如对配置文件的处理

        |--data //使用的数据结构，如 DMatrix

        |--gbm //定义了若分类器，如 gbtree 和 gblinear

        |--metric //定义评价函数

        |--objective //定义目标函数

        |--tree  //对树的一些列操作

从目录结构的角度看，代码层次结构，代码模块是非常干净利落的。

由于我们的重点不在于XGB的单机代码，如果对XGB感兴趣的，可以沿着cli_main.cc的执行入口，再到训练的核心方法CLITrain(param)，再到Learner::UpdateOneIter()的实际树更新逻辑，再到Learner里头实现ObjFunction::GetGradient()梯度求解的过程（包括Loss函数，一阶和二阶的导师计算）。

然后在Tr