Merge pull request #23 from i8run/master

Support MKL2017 DNN API With New Branch

.gitignore

@@ -20,3 +20,4 @@ project/plugins/project/
 
 # other
 *.txt
+*.swp # vim swap file

dl/src/main/scala/com/intel/analytics/sparkdl/nn/Container.scala

@@ -19,6 +19,7 @@ package com.intel.analytics.sparkdl.nn
 
 import com.intel.analytics.sparkdl.tensor.Tensor
 import com.intel.analytics.sparkdl.tensor.TensorNumericMath.TensorNumeric
+import com.intel.analytics.sparkdl.mkl.MKL
 
 import scala.collection.mutable.ArrayBuffer
 import scala.reflect.ClassTag
@@ -93,4 +94,25 @@ private[nn] abstract class Container[@specialized(Float, Double) T: ClassTag](
     })
     (result, offset, newIndexes)
   }
+
+  override def initMkl() : Unit = {
+    def containMkl(module : Module[T]) : Boolean = {
+      return if (module.toString.startsWith("mkl.")) true else false
+    }
+
+    for (i <- 0 until modules.length) {
+      if (containMkl(modules(i))) {
+        if (i >= 1 && containMkl(modules(i - 1))) {
+          ev.getType() match {
+            case "Float" => MKL.SetPrevFloat(modules(i - 1).getClassPtr(),
+                                             modules(i).getClassPtr())
+            case "Double" => MKL.SetPrevDouble(modules(i - 1).getClassPtr(),
+                                               modules(i).getClassPtr())
+          }
+        }
+      } else {
+        modules(i).initMkl()
+      }
+    }
+  }
 }

dl/src/main/scala/com/intel/analytics/sparkdl/nn/Module.scala

@@ -43,6 +43,17 @@ abstract class Module[T: ClassTag](implicit ev: TensorNumeric[T]) extends Serial
     if (this.name == null) this.toString else this.name
   }
 
+  private var needComputeBack = true
+
+  def setNeedComputeBack(need: Boolean): this.type = {
+    needComputeBack = need
+    this
+  }
+
+  def isNeedComputeBack(): Boolean = {
+    needComputeBack
+  }
+
   // list of sub modules
   val modules: ArrayBuffer[Module[T]] = ArrayBuffer[Module[T]]()
 
@@ -199,6 +210,10 @@ abstract class Module[T: ClassTag](implicit ev: TensorNumeric[T]) extends Serial
   def cloneModule(): Module[T] = {
     SerializationUtils.clone(this)
   }
+
+  // Support for mkl init.
+  def getClassPtr() : Long = {0L}
+  def initMkl() : Unit = {}
 }
 
 object Module {

dl/src/main/scala/com/intel/analytics/sparkdl/nn/mkl/BatchNormalization.scala

@@ -0,0 +1,220 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.intel.analytics.sparkdl.nn.mkl
+
+import com.intel.analytics.sparkdl.tensor.TensorNumericMath.TensorNumeric
+import com.intel.analytics.sparkdl.tensor.Tensor
+import com.intel.analytics.sparkdl.utils.RandomGenerator._
+import com.intel.analytics.sparkdl.nn.Module
+import com.intel.analytics.sparkdl.mkl.MKL
+
+import scala.language.implicitConversions
+
+import scala.reflect.ClassTag
+
+class SpatialBatchNormalization[@specialized(Float, Double) T: ClassTag](
+    val nOutput: Int,
+    val eps: Double = 1e-5,
+    val momentum: Double = 0.1,
+    val affine: Boolean = true)(implicit ev: TensorNumeric[T])
+    extends Module[T] {
+
+  require(nOutput > 0,
+          "To set affine=false call SpatialBatchNormalization(nFeature,  eps, momentum, false)")
+
+  val nDim = 2
+  val runningMean = Tensor[T](nOutput)
+  val runningVar = Tensor[T](nOutput).fill(ev.fromType[Int](1))
+  val saveMean = Tensor[T](nOutput)
+  val saveStd = Tensor[T](nOutput).fill(ev.fromType[Int](1))
+
+  private var classPtr = 0L
+  private var firstPass = true
+
+  override def getClassPtr(): Long = classPtr
+
+  val weight: Tensor[T] = if (affine) Tensor[T](nOutput) else null
+  val bias: Tensor[T] = if (affine) Tensor[T](nOutput) else null
+  gradWeight = if (affine) Tensor[T](nOutput) else null
+  gradBias = if (affine) Tensor[T](nOutput) else null
+
+  val useWeight: Boolean = if (weight != null) true else false
+  val useBias: Boolean = if (bias != null) true else false
+
+  if (affine) {
+    reset()
+  }
+
+  override def reset(): Unit = {
+    if (null != weight) {
+      weight.apply1(_ => ev.fromType[Double](RNG.uniform(0, 1)))
+    }
+
+    if (null != bias) {
+      bias.fill(ev.fromType[Int](0))
+    }
+
+    runningMean.zero()
+    runningVar.fill(ev.fromType[Int](1))
+  }
+
+  def checkInputDim(input: Tensor[T]): Unit = {
+    require(input.dim() == nDim,
+            s"only mini-batch supported (${nDim}D tensor), got ${input.dim()}D tensor instead")
+    require(input.size(2) == runningMean.nElement(),
+            s"got ${input.size(2)}-feature tensor, expected ${runningMean.nElement()}")
+  }
+
+  override def updateOutput(input: Tensor[T]): Tensor[T] = {
+    output.resizeAs(input)
+
+    val inputOffset = input.storageOffset() - 1;
+    val outputOffset = output.storageOffset() - 1;
+
+    val inputWidth = input.size(input.dim())
+    val inputHeight = input.size(input.dim() - 1)
+    val inputChannel = if (input.dim() <= 2) 1 else input.size(input.dim() - 2)
+    val inputNumber = if (input.dim() <= 3) 1 else input.size(input.dim() - 3)
+    // TODO we may set input.size(input.dim() - 3) == 1 if input.dim() == 3
+
+    val kernelOffset = weight.storageOffset() - 1
+    val biasOffset = bias.storageOffset() - 1
+
+    implicit def bool2int(b: Boolean) = if (b) 1 else 0
+    if (firstPass) {
+      ev.getType() match {
+        case "Float" =>
+          classPtr = MKL.BatchNormInitFloat(inputNumber,
+                                            inputChannel,
+                                            inputHeight,
+                                            inputWidth,
+                                            eps,
+                                            useWeight,
+                                            useBias,
+                                            4)
+        case "Double" =>
+          classPtr = MKL.BatchNormInitDouble(inputNumber,
+                                             inputChannel,
+                                             inputHeight,
+                                             inputWidth,
+                                             eps,
+                                             useBias,
+                                             useBias,
+                                             4)
+        case _ =>
+          throw new UnsupportedOperationException(s"Only Float/Double supported")
+      }
+      firstPass = false
+    }
+
+    ev.getType() match {
+      case "Float" =>
+        MKL.BatchNormForwardFloat(input.storage().array().asInstanceOf[Array[Float]],
+                                  inputOffset,
+                                  output.storage().array().asInstanceOf[Array[Float]],
+                                  outputOffset,
+                                  weight.storage().array().asInstanceOf[Array[Float]],
+                                  kernelOffset,
+                                  bias.storage().array().asInstanceOf[Array[Float]],
+                                  biasOffset,
+                                  classPtr)
+      case "Double" =>
+        MKL.BatchNormForwardDouble(input.storage().array().asInstanceOf[Array[Double]],
+                                   inputOffset,
+                                   output.storage().array().asInstanceOf[Array[Double]],
+                                   outputOffset,
+                                   weight.storage().array().asInstanceOf[Array[Double]],
+                                   kernelOffset,
+                                   bias.storage().array().asInstanceOf[Array[Double]],
+                                   biasOffset,
+                                   classPtr)
+      case _ =>
+        throw new UnsupportedOperationException(s"Only Float/Double supported")
+    }
+    output
+  }
+
+  override def updateGradInput(input: Tensor[T], gradOutput: Tensor[T]): Tensor[T] = {
+    gradInput.resizeAs(input)
+
+    val inputOffset = input.storageOffset() - 1;
+    val outputOffset = output.storageOffset() - 1;
+
+    val inputWidth = input.size(input.dim())
+    val inputHeight = input.size(input.dim() - 1)
+    val inputChannel = if (input.dim() <= 2) 1 else input.size(input.dim() - 2)
+    val inputNumber = if (input.dim() <= 3) 1 else input.size(input.dim() - 3)
+    // TODO we may set input.size(input.dim() - 3) == 1 if input.dim() == 3
+
+    val kernelOffset = weight.storageOffset() - 1
+    val biasOffset = bias.storageOffset() - 1
+
+    val kernelDiffOffset = gradWeight.storageOffset() - 1
+    val biasDiffOffset = gradBias.storageOffset() - 1
+
+    val gradOutputOffset = gradOutput.storageOffset() - 1
+    val gradInputOffset = gradInput.storageOffset() - 1
+
+    implicit def bool2int(b: Boolean) = if (b) 1 else 0
+    ev.getType() match {
+      case "Float" =>
+        MKL.BatchNormBackwardFloat(input.storage().array().asInstanceOf[Array[Float]],
+                                   inputOffset,
+                                   gradOutput.storage().array().asInstanceOf[Array[Float]],
+                                   gradOutputOffset,
+                                   gradInput.storage().array().asInstanceOf[Array[Float]],
+                                   gradInputOffset,
+                                   gradWeight.storage().array().asInstanceOf[Array[Float]],
+                                   kernelDiffOffset,
+                                   gradBias.storage().array().asInstanceOf[Array[Float]],
+                                   biasDiffOffset,
+                                   classPtr)
+      case "Double" =>
+        MKL.BatchNormBackwardDouble(input.storage().array().asInstanceOf[Array[Double]],
+                                    inputOffset,
+                                    gradOutput.storage().array().asInstanceOf[Array[Double]],
+                                    gradOutputOffset,
+                                    gradInput.storage().array().asInstanceOf[Array[Double]],
+                                    gradInputOffset,
+                                    gradWeight.storage().array().asInstanceOf[Array[Double]],
+                                    kernelDiffOffset,
+                                    gradBias.storage().array().asInstanceOf[Array[Double]],
+                                    biasDiffOffset,
+                                    classPtr)
+      case _ =>
+        throw new UnsupportedOperationException(s"Only Float/Double supported")
+    }
+
+    gradInput
+  }
+
+  override def accGradParameters(input: Tensor[T], gradOutput: Tensor[T], scale: Double): Unit = {}
+
+  override def zeroGradParameters(): Unit = {
+    gradWeight.zero()
+    gradBias.zero()
+  }
+
+  override def parameters(): (Array[Tensor[T]], Array[Tensor[T]]) = {
+    (Array(this.weight, this.bias), Array(this.gradWeight, this.gradBias))
+  }
+
+  override def toString(): String = {
+    s"mkl.BatchNormalization[${ev.getType()}]($nOutput, $eps, $momentum, $affine)"
+  }
+}