Refactoring. TopologyModel trait as a common model interface

avulanov · May 15, 2015 · cb03fe0 · cb03fe0
1 parent 7f15956
commit cb03fe0
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 8 deletions.
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/ann/Layer.scala b/mllib/src/main/scala/org/apache/spark/mllib/ann/Layer.scala
@@ -453,7 +453,7 @@ object FeedForwardModel {
 
 /* Neural network gradient. Does nothing but calling Model's gradient
 * */
-class ANNGradient(topology: FeedForwardTopology, dataStacker: DataStacker) extends Gradient {
+class ANNGradient(topology: Topology, dataStacker: DataStacker) extends Gradient {
 
   override def compute(data: Vector, label: Double, weights: Vector): (Vector, Double) = {
     val gradient = Vectors.zeros(weights.size)
@@ -464,7 +464,6 @@ class ANNGradient(topology: FeedForwardTopology, dataStacker: DataStacker) exten
   override def compute(data: Vector, label: Double, weights: Vector,
                        cumGradient: Vector): Double = {
     val (input, target, realBatchSize) = dataStacker.unstack(data)
-    //val model = FeedForwardModel(topology, weights)
     val model = topology.getInstance(weights)
     model.computeGradient(input, target, cumGradient, realBatchSize)
   }
@@ -528,11 +527,11 @@ private class ANNUpdater extends Updater {
 }
 /* MLlib-style trainer class that trains a network given the data and topology
 * */
-class FeedForwardTrainer (topology: FeedForwardTopology, val inputSize: Int,
+class FeedForwardTrainer (topology: Topology, val inputSize: Int,
                           val outputSize: Int) extends Serializable {
 
   // TODO: what if we need to pass random seed?
-  private var _weights = FeedForwardModel(topology).weights()
+  private var _weights = topology.getInstance(11L).weights()//FeedForwardModel(topology).weights()
   private var _batchSize = 1
   private var dataStacker = new DataStacker(_batchSize, inputSize, outputSize)
   private var _gradient: Gradient = new ANNGradient(topology, dataStacker)
@@ -595,9 +594,10 @@ class FeedForwardTrainer (topology: FeedForwardTopology, val inputSize: Int,
     }
   }
 
-  def train(data: RDD[(Vector, Vector)]): FeedForwardModel = {
+  def train(data: RDD[(Vector, Vector)]): TopologyModel = {
     val newWeights = optimizer.optimize(dataStacker.stack(data), getWeights)
-    FeedForwardModel(topology, newWeights)
+    //FeedForwardModel(topology, newWeights)
+    topology.getInstance(newWeights)
   }
 
 }
diff --git a/mllib/src/main/scala/org/apache/spark/mllib/classification/ANNClassifier.scala b/mllib/src/main/scala/org/apache/spark/mllib/classification/ANNClassifier.scala
@@ -17,7 +17,7 @@
 
 package org.apache.spark.mllib.classification
 
-import org.apache.spark.mllib.ann.{FeedForwardTrainer, FeedForwardModel}
+import org.apache.spark.mllib.ann.{TopologyModel, Topology, FeedForwardTrainer, FeedForwardModel}
 import org.apache.spark.mllib.linalg.Vector
 import org.apache.spark.mllib.regression.LabeledPoint
 import org.apache.spark.rdd.RDD
@@ -37,7 +37,7 @@ object LabelConverter {
   }
 }
 
-class ANNClassifierModel (val annModel: FeedForwardModel)
+class ANNClassifierModel (val annModel: TopologyModel)
   extends ClassificationModel with Serializable {
   /**
    * Predict values for the given data set using the model trained.