remove unnecessary allocations from caffenet

Author: robertnishihara

models/test/test.prototxt

@@ -0,0 +1,13 @@
+name: "test"
+input: "data"
+input_shape {
+  dim: 256
+  dim: 3
+  dim: 227
+  dim: 227
+}
+input: "label"
+input_shape {
+  dim: 256
+  dim: 1
+}

src/main/scala/apps/CifarApp.scala

@@ -111,7 +111,7 @@ object CifarApp {
             val numTestBatches = workerStore.get[Int]("testPartitionSize") / testBatchSize
             var accuracy = 0F
             for (j <- 0 to numTestBatches - 1) {
-              val out = workerStore.get[CaffeSolver]("solver").trainNet.forward(testIt)
+              val out = workerStore.get[CaffeSolver]("solver").trainNet.forward(testIt, List("accuracy", "loss", "prob"))
               accuracy += out("accuracy").get(Array())
             }
             Array[(Float, Int)]((accuracy, numTestBatches)).iterator

src/main/scala/apps/ImageNetApp.scala

@@ -111,7 +111,7 @@ object ImageNetApp {
             val numTestBatches = workerStore.get[Int]("testPartitionSize") / testBatchSize
             var accuracy = 0F
             for (j <- 0 to numTestBatches - 1) {
-              val out = workerStore.get[CaffeSolver]("solver").trainNet.forward(testIt)
+              val out = workerStore.get[CaffeSolver]("solver").trainNet.forward(testIt, List("accuracy"))
               accuracy += out("accuracy").get(Array())
             }
             Array[(Float, Int)]((accuracy, numTestBatches)).iterator

src/main/scala/libs/CaffeNet.scala

@@ -26,19 +26,19 @@ object CaffeNet {
 }
 
 class CaffeNet(netParam: NetParameter, schema: StructType, preprocessor: Preprocessor, caffeNet: FloatNet) {
-  private val inputSize = netParam.input_size
-  private val batchSize = netParam.input_shape(0).dim(0).toInt
-  private val transformations = new Array[Any => NDArray](inputSize)
+  val inputSize = netParam.input_size
+  val batchSize = netParam.input_shape(0).dim(0).toInt
+  private val transformations = new Array[(Any, Array[Float]) => Unit](inputSize)
   private val inputIndices = new Array[Int](inputSize)
   private val columnNames = schema.map(entry => entry.name)
   // private val caffeNet = new FloatNet(netParam)
   private val inputRef = new Array[FloatBlob](inputSize)
   def getNet = caffeNet // TODO: For debugging
 
-  private val numOutputs = caffeNet.num_outputs
-  private val numLayers = caffeNet.layers().size.toInt
-  private val layerNames = List.range(0, numLayers).map(i => caffeNet.layers.get(i).layer_param.name.getString)
-  private val numLayerBlobs = List.range(0, numLayers).map(i => caffeNet.layers.get(i).blobs().size.toInt)
+  val numOutputs = caffeNet.num_outputs
+  val numLayers = caffeNet.layers().size.toInt
+  val layerNames = List.range(0, numLayers).map(i => caffeNet.layers.get(i).layer_param.name.getString)
+  val numLayerBlobs = List.range(0, numLayers).map(i => caffeNet.layers.get(i).blobs().size.toInt)
 
   for (i <- 0 to inputSize - 1) {
     val name = netParam.input(i).getString
@@ -58,43 +58,35 @@ class CaffeNet(netParam: NetParameter, schema: StructType, preprocessor: Preproc
     inputRef(i) = new FloatBlob(dims)
     inputs.put(i, inputRef(i))
   }
-  val inputBuffer = new Array[Array[Float]](inputSize)
+  // in `inputBuffer`, the first index indexes the input argument, the second
+  // index indexes into the batch, the third index indexes the values in the
+  // data
+  val inputBuffer = new Array[Array[Array[Float]]](inputSize)
   val inputBufferSize = new Array[Int](inputSize)
   for (i <- 0 to inputSize - 1) {
     inputBufferSize(i) = JavaCPPUtils.getInputShape(netParam, i).drop(1).product // drop 1 to ignore batchSize
-    inputBuffer(i) = new Array[Float](inputBufferSize(i) * batchSize)
+    inputBuffer(i) = new Array[Array[Float]](batchSize)
+    for (batchIndex <- 0 to batchSize - 1) {
+      inputBuffer(i)(batchIndex) = new Array[Float](inputBufferSize(i))
+    }
   }
 
-  def transformInto(iterator: Iterator[Row], data: FloatBlobVector) = {
+  def transformInto(iterator: Iterator[Row], inputs: FloatBlobVector) = {
     var batchIndex = 0
     while (iterator.hasNext && batchIndex != batchSize) {
       val row = iterator.next
       for (i <- 0 to inputSize - 1) {
-        val result = transformations(i)(row(inputIndices(i)))
-        val flatArray = result.toFlat() // TODO: Make this efficient
-        System.arraycopy(flatArray, 0, inputBuffer(i), batchIndex * inputBufferSize(i), inputBufferSize(i))
+        transformations(i)(row(inputIndices(i)), inputBuffer(i)(batchIndex))
       }
       batchIndex += 1
     }
-    for (i <- 0 to inputSize - 1) {
-      val blob = data.get(i)
-      val buffer = blob.mutable_cpu_data()
-      buffer.put(inputBuffer(i), 0, batchSize * inputBufferSize(i))
-    }
+    JavaCPPUtils.arraysToFloatBlobVector(inputBuffer, inputs, batchSize, inputBufferSize, inputSize)
   }
 
   def forward(rowIt: Iterator[Row], dataBlobNames: List[String] = List[String]()): Map[String, NDArray] = {
     transformInto(rowIt, inputs)
-    val tops = caffeNet.Forward(inputs)
+    caffeNet.Forward(inputs)
     val outputs = Map[String, NDArray]()
-    for (j <- 0 to numOutputs - 1) {
-      val outputName = caffeNet.blob_names().get(caffeNet.output_blob_indices().get(j)).getString
-      val top = tops.get(j)
-      val shape = Array.range(0, top.num_axes).map(i => top.shape.get(i))
-      val output = new Array[Float](shape.product)
-      top.cpu_data().get(output, 0, shape.product)
-      outputs += (outputName -> NDArray(output, shape))
-    }
     for (name <- dataBlobNames) {
       val floatBlob = caffeNet.blob_by_name(name)
       if (floatBlob == null) {

src/main/scala/libs/JavaCPPUtils.scala

@@ -25,7 +25,7 @@ object JavaCPPUtils {
     shape
   }
 
- def getInputShape(netParam: NetParameter, i: Int): Array[Int] = {
+  def getInputShape(netParam: NetParameter, i: Int): Array[Int] = {
     val numAxes = netParam.input_shape(i).dim_size
     val shape = new Array[Int](numAxes)
     for (j <- 0 to numAxes - 1) {
@@ -34,5 +34,41 @@ object JavaCPPUtils {
     shape
   }
 
+  def arraysToFloatBlobVector(inputBuffer: Array[Array[Array[Float]]], inputs: FloatBlobVector, batchSize: Int, inputBufferSize: Array[Int], inputSize: Int) = {
+    for (i <- 0 to inputSize - 1) {
+      val blob = inputs.get(i)
+      val buffer = blob.mutable_cpu_data()
+      var batchIndex = 0
+      while (batchIndex < batchSize) {
+        var j = 0
+        while (j < inputBufferSize(i)) {
+          // it'd be preferable to do this with one call, but JavaCPP's FloatPointer API has confusing semantics
+          buffer.put(inputBufferSize(i) * batchIndex + j, inputBuffer(i)(batchIndex)(j))
+          j += 1
+        }
+        batchIndex += 1
+      }
+    }
+  }
+
+  // this method is just for testing
+  def arraysFromFloatBlobVector(inputs: FloatBlobVector, batchSize: Int, inputBufferSize: Array[Int], inputSize: Int): Array[Array[Array[Float]]] = {
+    val result = new Array[Array[Array[Float]]](inputSize)
+    for (i <- 0 to inputSize - 1) {
+      result(i) = new Array[Array[Float]](batchSize)
+      val blob = inputs.get(i)
+      val buffer = blob.cpu_data()
+      for (batchIndex <- 0 to batchSize - 1) {
+        result(i)(batchIndex) = new Array[Float](inputBufferSize(i))
+        var j = 0
+        while (j < inputBufferSize(i)) {
+          // it'd be preferable to do this with one call, but JavaCPP's FloatPointer API has confusing semantics
+          result(i)(batchIndex)(j) = buffer.get(inputBufferSize(i) * batchIndex + j)
+          j += 1
+        }
+      }
+    }
+    return result
+  }
 
 }

src/main/scala/libs/Preprocessor.scala

@@ -6,8 +6,10 @@ import org.apache.spark.sql.types._
 import org.apache.spark.sql.{DataFrame, Row}
 import scala.collection.mutable._
 
+// The Preprocessor is provides a function for reading data from a dataframe row
+// into the net
 trait Preprocessor {
-  def convert(name: String, shape: Array[Int]): Any => NDArray
+  def convert(name: String, shape: Array[Int]): (Any, Array[Float]) => Unit
 }
 
 trait TensorFlowPreprocessor {
@@ -20,62 +22,68 @@ trait TensorFlowPreprocessor {
 // allocation. This is designed to be easier to understand, whereas the
 // ImageNetPreprocessor is designed to be faster.
 class DefaultPreprocessor(schema: StructType) extends Preprocessor {
-  def convert(name: String, shape: Array[Int]): Any => NDArray = {
+  def convert(name: String, shape: Array[Int]): (Any, Array[Float]) => Unit = {
     schema(name).dataType match {
-      case FloatType => (element: Any) => {
-        NDArray(Array[Float](element.asInstanceOf[Float]), shape)
+      case FloatType => (element: Any, buffer: Array[Float]) => {
+        if (buffer.length != shape.product) { throw new Exception("buffer.length and shape.product don't agree, buffer has length " + buffer.length.toString + ", but shape is " + shape.deep.toString) }
+        NDArray(Array[Float](element.asInstanceOf[Float]), shape).flatCopy(buffer)
       }
-      case DoubleType => (element: Any) => {
-        NDArray(Array[Float](element.asInstanceOf[Double].toFloat), shape)
+      case DoubleType => (element: Any, buffer: Array[Float]) => {
+        if (buffer.length != shape.product) { throw new Exception("buffer.length and shape.product don't agree, buffer has length " + buffer.length.toString + ", but shape is " + shape.deep.toString) }
+        NDArray(Array[Float](element.asInstanceOf[Double].toFloat), shape).flatCopy(buffer)
       }
-      case IntegerType => (element: Any) => {
-        NDArray(Array[Float](element.asInstanceOf[Int].toFloat), shape)
+      case IntegerType => (element: Any, buffer: Array[Float]) => {
+        if (buffer.length != shape.product) { throw new Exception("buffer.length and shape.product don't agree, buffer has length " + buffer.length.toString + ", but shape is " + shape.deep.toString) }
+        NDArray(Array[Float](element.asInstanceOf[Int].toFloat), shape).flatCopy(buffer)
       }
-      case LongType => (element: Any) => {
-        NDArray(Array[Float](element.asInstanceOf[Long].toFloat), shape)
+      case LongType => (element: Any, buffer: Array[Float]) => {
+        if (buffer.length != shape.product) { throw new Exception("buffer.length and shape.product don't agree, buffer has length " + buffer.length.toString + ", but shape is " + shape.deep.toString) }
+        NDArray(Array[Float](element.asInstanceOf[Long].toFloat), shape).flatCopy(buffer)
       }
-      case BinaryType => (element: Any) => {
-        NDArray(element.asInstanceOf[Array[Byte]].map(e => (e & 0xFF).toFloat), shape)
+      case BinaryType => (element: Any, buffer: Array[Float]) => {
+        if (buffer.length != shape.product) { throw new Exception("buffer.length and shape.product don't agree, buffer has length " + buffer.length.toString + ", but shape is " + shape.deep.toString) }
+        NDArray(element.asInstanceOf[Array[Byte]].map(e => (e & 0xFF).toFloat), shape).flatCopy(buffer)
       }
-      // case ArrayType(IntegerType, true) => (element: Any) => {} // TODO(rkn): implement
-      case ArrayType(FloatType, true) => (element: Any) => {
+      case ArrayType(FloatType, true) => (element: Any, buffer: Array[Float]) => {
+        if (buffer.length != shape.product) { throw new Exception("buffer.length and shape.product don't agree, buffer has length " + buffer.length.toString + ", but shape is " + shape.deep.toString) }
         element match {
-          case element: Array[Float] => NDArray(element.asInstanceOf[Array[Float]], shape)
-          case element: WrappedArray[Float] => NDArray(element.asInstanceOf[WrappedArray[Float]].toArray, shape)
-          case element: ArrayBuffer[Float] => NDArray(element.asInstanceOf[ArrayBuffer[Float]].toArray, shape)
+          case element: Array[Float] => NDArray(element.asInstanceOf[Array[Float]], shape).flatCopy(buffer)
+          case element: WrappedArray[Float] => NDArray(element.asInstanceOf[WrappedArray[Float]].toArray, shape).flatCopy(buffer)
+          case element: ArrayBuffer[Float] => NDArray(element.asInstanceOf[ArrayBuffer[Float]].toArray, shape).flatCopy(buffer)
         }
       }
-      // case ArrayType(DoubleType, true) => (element: Any) => {} // TODO(rkn): implement
-      // case ArrayType(LongType, true) => (element: Any) => {} // TODO(rkn): implement
     }
   }
 }
 
 class ImageNetPreprocessor(schema: StructType, meanImage: Array[Float], fullHeight: Int = 256, fullWidth: Int = 256, croppedHeight: Int = 227, croppedWidth: Int = 227) extends Preprocessor {
-  def convert(name: String, shape: Array[Int]): Any => NDArray = {
+  def convert(name: String, shape: Array[Int]): (Any, Array[Float]) => Unit = {
     schema(name).dataType match {
-      case IntegerType => (element: Any) => {
-        NDArray(Array[Float](element.asInstanceOf[Int].toFloat), shape)
+      case IntegerType => (element: Any, buffer: Array[Float]) => {
+        if (buffer.length != shape.product) { throw new Exception("buffer.length and shape.product don't agree, buffer has length " + buffer.length.toString + ", but shape is " + shape.deep.toString) }
+        NDArray(Array[Float](element.asInstanceOf[Int].toFloat), shape).flatCopy(buffer)
       }
       case BinaryType => {
         if (shape(0) != 3) {
           throw new IllegalArgumentException("Expecting input image to have 3 channels.")
         }
-        val buffer = new Array[Float](3 * fullHeight * fullWidth)
-        (element: Any) => {
+        val tempBuffer = new Array[Float](3 * fullHeight * fullWidth)
+        (element: Any, buffer: Array[Float]) => {
+          if (buffer.length != shape.product) { throw new Exception("buffer.length and shape.product don't agree, buffer has length " + buffer.length.toString + ", but shape is " + shape.deep.toString) }
           element match {
             case element: Array[Byte] => {
               var index = 0
               while (index < 3 * fullHeight * fullWidth) {
-                buffer(index) = (element(index) & 0xFF).toFloat - meanImage(index)
+                tempBuffer(index) = (element(index) & 0xFF).toFloat - meanImage(index)
                 index += 1
               }
             }
           }
           val heightOffset = Random.nextInt(fullHeight - croppedHeight + 1)
           val widthOffset = Random.nextInt(fullWidth - croppedWidth + 1)
-          NDArray(buffer.clone, Array[Int](shape(0), fullHeight, fullWidth)).subarray(Array[Int](0, heightOffset, widthOffset), Array[Int](shape(0), heightOffset + croppedHeight, widthOffset + croppedWidth))
-          // TODO(rkn): probably don't want to call buffer.clone
+          val lowerIndices = Array[Int](0, heightOffset, widthOffset)
+          val upperIndices = Array[Int](shape(0), heightOffset + croppedHeight, widthOffset + croppedWidth)
+          NDArray(tempBuffer, Array[Int](shape(0), fullHeight, fullWidth)).subarray(lowerIndices, upperIndices).flatCopy(buffer)
         }
       }
     }

src/test/scala/apps/LoadAdultDataSpec.scala

@@ -19,11 +19,13 @@ class LoadAdultDataSpec extends FlatSpec {
 
     val function0 = preprocessor.convert("C0", Array[Int](1))
     val function2 = preprocessor.convert("C2", Array[Int](1))
-    val result0 = function0(df.take(1)(0)(0))
-    val result2 = function2(df.take(1)(0)(2))
+    val result0 = new Array[Float](1)
+    val result2 = new Array[Float](1)
+    function0(df.take(1)(0)(0), result0)
+    function2(df.take(1)(0)(2), result2)
 
-    assert((result0.get(Array[Int](0)) - 39.0).abs <= 1e-4)
-    assert((result2.get(Array[Int](0)) - 77516.0).abs <= 1e-4)
+    assert((result0(0) - 39.0).abs <= 1e-4)
+    assert((result2(0) - 77516.0).abs <= 1e-4)
 
     sc.stop()
   }