support kryo serialization

This commit is contained in:
CodingCat
2016-03-13 11:55:14 -04:00
parent 9011acf52b
commit f2ef958ebb
6 changed files with 88 additions and 10 deletions

View File

@@ -18,12 +18,15 @@ package ml.dmlc.xgboost4j.scala
import java.io.IOException
import com.esotericsoftware.kryo.io.{Output, Input}
import com.esotericsoftware.kryo.{Kryo, KryoSerializable}
import ml.dmlc.xgboost4j.java.{Booster => JBooster}
import ml.dmlc.xgboost4j.java.XGBoostError
import scala.collection.JavaConverters._
import scala.collection.mutable
class Booster private[xgboost4j](booster: JBooster) extends Serializable {
class Booster private[xgboost4j](private var booster: JBooster)
extends Serializable with KryoSerializable {
/**
* Set parameter to the Booster.
@@ -193,4 +196,12 @@ class Booster private[xgboost4j](booster: JBooster) extends Serializable {
super.finalize()
dispose
}
override def write(kryo: Kryo, output: Output): Unit = {
kryo.writeObject(output, booster)
}
override def read(kryo: Kryo, input: Input): Unit = {
booster = kryo.readObject(input, classOf[JBooster])
}
}