spark mlib java_如何在Java应用里集成Spark MLlib训练好的模型做预测

前言

昨天媛媛说，你是不是很久没写博客了。我说上一篇1.26号，昨天3.26号，刚好两个月，心中也略微有些愧疚。今天正好有个好朋友问，怎么在Java应用里集成Spark MLlib训练好的模型。在StreamingPro里其实都有实际的使用例子，但是如果有一篇文章讲述下，我觉得应该能让更多人获得帮助

追本溯源

记得我之前吐槽过Spark MLlib的设计，也是因为一个朋友使用了spark MLlib的pipeline做训练，然后他把这个pipeline放到了spring boot里，结果做预测的时候奇慢无比，一条记录inference需要30多秒。为什么会这么慢呢？原因是Spark MLlib 是以批处理为核心设计理念的。比如上面朋友遇到的坑是有一部分原因来源于word2vec的transform方法：

@Since("2.0.0")

override def transform(dataset: Dataset[_]): DataFrame = {

transformSchema(dataset.schema, logging = true)

val vectors = wordVectors.getVectors

.mapValues(vv => Vectors.dense(vv.map(_.toDouble)))

.map(identity) // mapValues doesn't return a serializable map (SI-7005)

val bVectors = dataset.sparkSession.sparkContext.broadcast(vectors)

val d = $(vectorSize)

来一条数据(通常API应用都是如此)，他需要先获得vectors(词到vector的映射)对象，假设你有十万个词，

def getVectors: Map[String, Array[Float]] = {

wordIndex.map { case (word, ind) =>

(word, wordVectors.slice(vectorSize * ind, vectorSize * ind + vectorSize))

}

}

每次请求他都要做如上调用和计算。接着还需要把这些东西(这个可能就比较大了，几百M或者几个G都有可能)广播出去。

所以注定快不了。

把model集成到Java 服务里实例

假设你使用贝叶斯训练了一个模型，你需要保存下这个模型，保存的方式如下：

val nb = new NaiveBayes()

//做些参数配置和训练过程

.....

//保存模型

nb.write.overwrite().save(path + "/" + modelIndex)

接着，在你的Java/scala程序里，引入spark core，spark mllib等包。加载模型：

val model = NaiveBayesModel.load(tempPath)

这个时候因为要做预测，我们为了性能，不能直接调用model的transform方法，你仔细观察发现，我们需要通过反射调用两个方法，就能实现分类。第一个是predictRaw方法，该方法输入一个向量，输出也为一个向量。我们其实不需要向量，我们需要的是一个分类的id。predictRaw 方法在model里，但是没办法直接调用，因为是私有的：

override protected def predictRaw(features: Vector): Vector = {

$(modelType) match {

case Multinomial =>

multinomialCalculation(features)

case Bernoulli =>

bernoulliCalculation(features)

case _ =>

// This should never happen.

throw new UnknownError(s"Invalid modelType: ${$(modelType)}.")

}

}

所以我们需要通过反射来完成：

val predictRaw = model.getClass.getMethod("predictRaw", classOf[Vector]).invoke(model, vec).asInstanceOf[Vector]

现在我们已经得到了predctRaw的结果，接着我们要用raw2probability 把向量转化为一个概率分布，因为spark 版本不同，该方法的签名也略有变化，所以可能要做下版本适配：

val raw2probabilityMethod = if (sparkSession.version.startsWith("2.3")) "raw2probabilityInPlace" else "raw2probability"

val raw2probability = model.getClass.getMethod(raw2probabilityMethod, classOf[Vector]).invoke(model, predictRaw).asInstanceOf[Vector]

raw2probability 其实也还是一个向量，这个向量的长度是分类的数目，每个位置的值是概率。所以所以我们只要拿到最大的那个概率值所在的位置就行：

val categoryId = raw2probability.argmax

这个时候categoryId 就是我们预测的分类了。

截止到目前我们已经完成了作为一个普通java/scala 方法的调用流程。如果我不想用在应用程序里，而是放到spark 流式计算里呢？或者批处理也行，那么这个时候你只需要封装一个UDF函数即可：

val models = sparkSession.sparkContext.broadcast(_model.asInstanceOf[ArrayBuffer[NaiveBayesModel]])

val f2 = (vec: Vector) => {

models.value.map { model =>

val predictRaw = model.getClass.getMethod("predictRaw", classOf[Vector]).invoke(model, vec).asInstanceOf[Vector]

val raw2probability = model.getClass.getMethod(raw2probabilityMethod, classOf[Vector]).invoke(model, predictRaw).asInstanceOf[Vector]

//model.getClass.getMethod("probability2prediction", classOf[Vector]).invoke(model, raw2probability).asInstanceOf[Vector]

raw2probability

}

}

sparkSession.udf.register(name , f2)

上面的例子可以参考StreamingPro 中streaming.dsl.mmlib.algs.SQLNaiveBayes的代码。不同的算法因为内部实现不同，我们使用起来也会略微有些区别。

总结

Spark MLlib学习了SKLearn里的transform和fit的概念，但是因为设计上还是遵循批处理的方式，实际部署后会有很大的性能瓶颈，不适合那种数据一条一条过来需要快速响应的预测流程，所以需要调用一些内部的API来完成最后的预测。