patch on memory leaks fix

scala-package/examples/src/main/scala/org/apache/mxnetexamples/neuralstyle/ModelVgg19.scala

@@ -20,11 +20,11 @@ package org.apache.mxnetexamples.neuralstyle
 import org.apache.mxnet.{Context, Executor, NDArray, Shape, Symbol}
 
 /**
- * Definition for the neuralstyle network and initialize it with pretrained weight
- */
+  * Definition for the neuralstyle network and initialize it with pretrained weight
+  */
 object ModelVgg19 {
   case class ConvExecutor(executor: Executor, data: NDArray, dataGrad: NDArray,
-                      style: Array[NDArray], content: NDArray, argDict: Map[String, NDArray])
+                          style: Array[NDArray], content: NDArray, argDict: Map[String, NDArray])
 
   def ConvRelu(data : Symbol, convName : String, reluName : String,
                numFilter : Int, kernel : (Int, Int) = (3, 3),
@@ -33,7 +33,9 @@ object ModelVgg19 {
       pad = Some(Shape(1, 1)), kernel = Shape(kernel._1, kernel._2),
       stride = Some(Shape(stride._1, stride._2)), no_bias = Some(false),
       workspace = Some(1024), name = convName)
-    Symbol.api.relu(data = Some(conv), name = reluName)
+    val relu = Symbol.api.relu(data = Some(conv), name = reluName)
+    conv.dispose()
+    relu
   }
 
   def getSymbol: (Symbol, Symbol) = {
@@ -81,7 +83,7 @@ object ModelVgg19 {
   }
 
   def getExecutor(style: Symbol, content: Symbol, modelPath: String,
-      inputSize: (Int, Int), ctx: Context): ConvExecutor = {
+                  inputSize: (Int, Int), ctx: Context): ConvExecutor = {
     val out = Symbol.Group(style, content)
     // make executor
     val (argShapes, outputShapes, auxShapes) = out.inferShape(
@@ -95,15 +97,17 @@ object ModelVgg19 {
       val key = s"arg:$name"
       if (pretrained.contains(key)) argDict(name).set(pretrained(key))
     }
+    pretrained.foreach(ele => ele._2.dispose())
     val executor = out.bind(ctx, argDict, gradDict)
+    out.dispose()
     val outArray = executor.outputs
     ConvExecutor(executor = executor,
-                              data = argDict("data"),
-                              dataGrad = gradDict("data"),
-                              style = outArray.take(outArray.length - 1),
-                              content = outArray(outArray.length - 1),
-                              argDict = argDict)
-    }
+      data = argDict("data"),
+      dataGrad = gradDict("data"),
+      style = outArray.take(outArray.length - 1),
+      content = outArray(outArray.length - 1),
+      argDict = argDict)
+  }
 
   def getModel(modelPath: String, inputSize: (Int, Int), ctx: Context): ConvExecutor = {
     val (style, content) = getSymbol

scala-package/examples/src/main/scala/org/apache/mxnetexamples/neuralstyle/NeuralStyle.scala

@@ -28,10 +28,11 @@ import org.kohsuke.args4j.{CmdLineParser, Option}
 import org.slf4j.LoggerFactory
 
 import scala.collection.JavaConverters._
+import scala.collection.mutable.ListBuffer
 
 /**
- * An Implementation of the paper A Neural Algorithm of Artistic Style
- */
+  * An Implementation of the paper A Neural Algorithm of Artistic Style
+  */
 object NeuralStyle {
   case class NSExecutor(executor: Executor, data: NDArray, dataGrad: NDArray)
 
@@ -108,25 +109,32 @@ object NeuralStyle {
     var gradScale = List[Int]()
     for (i <- 0 until style.listOutputs().length) {
       val shape = outputShape(i)
-      val x = Symbol.Reshape()()(Map("data" -> style.get(i),
-          "target_shape" -> Shape(shape(1), shape(2) * shape(3))))
-      // use fully connected to quickly do dot(x, x^T)
-      val gram = Symbol.FullyConnected()()(Map("data" -> x, "weight" -> x,
-          "no_bias" -> true, "num_hidden" -> shape(1)))
+      val x = Symbol.api.Reshape(data = Some(style.get(i)),
+        target_shape = Some(Shape(shape(1), shape(2) * shape(3))))
+      val gram = Symbol.api.FullyConnected(data = Some(x), weight = Some(x),
+        no_bias = Some(true), num_hidden = shape(1))
+      x.dispose()
       gramList = gramList :+ gram
       gradScale = gradScale :+ (shape(1) * shape(2) * shape(3) * shape(1))
     }
     (Symbol.Group(gramList: _*), gradScale)
   }
 
   def getLoss(gram: Symbol, content: Symbol): (Symbol, Symbol) = {
-    var gramLoss = List[Symbol]()
+    var gramLoss = ListBuffer[Symbol]()
     for (i <- 0 until gram.listOutputs().length) {
       val gvar = Symbol.Variable(s"target_gram_$i")
-      gramLoss = gramLoss :+ Symbol.sum()(Symbol.square()(gvar - gram.get(i))())()
+      Symbol.api.square(data = Some(gvar - gram.get(i)))
+      gramLoss += Symbol.api.sum(
+        Some(Symbol.api.square(data = Some(gvar - gram.get(i))))
+      )
+      gvar.dispose()
     }
+    gram.dispose()
     val cvar = Symbol.Variable("target_content")
-    val contentLoss = Symbol.sum()(Symbol.square()(cvar - content)())()
+    val contentLoss = Symbol.api.sum(
+      Some(Symbol.api.square(Some(cvar - content)))
+    )
     (Symbol.Group(gramLoss: _*), contentLoss)
   }
 
@@ -137,12 +145,13 @@ object NeuralStyle {
     val nChannel = img.shape(1)
     val sImg = Symbol.Variable("img")
     val sKernel = Symbol.Variable("kernel")
-    val channels = Symbol.SliceChannel()(sImg)(Map("num_outputs" -> nChannel))
-    val out = Symbol.Concat()((0 until nChannel).map { i =>
-      Symbol.Convolution()()(Map("data" -> channels.get(i), "weight" -> sKernel,
-                    "num_filter" -> 1, "kernel" -> "(3,3)", "pad" -> "(1,1)",
-                    "no_bias" -> true, "stride" -> "(1,1)"))
-    }: _*)() * tvWeight
+    val channels = Symbol.api.SliceChannel(data = Some(sImg), num_outputs = nChannel)
+    val result = (0 until nChannel).map { i =>
+      Symbol.api.Convolution(data = Some(channels.get(i)), weight = Some(sKernel),
+        num_filter = 1, kernel = Shape(3, 3), pad = Some(Shape(1, 1)), no_bias = Some(true),
+        stride = Some(Shape(1, 1)))
+    }.toArray
+    val out = Symbol.api.Concat(result, result.length) * tvWeight
     val kernel = {
       val tmp = NDArray.empty(Shape(1, 1, 3, 3), ctx)
       tmp.set(Array[Float](0, -1, 0, -1, 4, -1, 0, -1, 0))
@@ -157,7 +166,7 @@ object NeuralStyle {
 
   //scalastyle:off
   def runTraining(model : String, contentImage : String, styleImage: String, dev : Context,
-                   modelPath : String, outputDir : String, styleWeight : Float,
+                  modelPath : String, outputDir : String, styleWeight : Float,
                   contentWeight : Float, tvWeight : Float, gaussianRadius : Int,
                   lr: Float, maxNumEpochs: Int, maxLongEdge: Int,
                   saveEpochs : Int, stopEps: Float) : Unit = {
@@ -179,6 +188,12 @@ object NeuralStyle {
     val contentArray = modelExecutor.content.copyTo(Context.cpu())
 
     // delete the executor
+    modelExecutor.argDict.foreach(ele => ele._2.dispose())
+    modelExecutor.content.dispose()
+    modelExecutor.data.dispose()
+    modelExecutor.dataGrad.dispose()
+    modelExecutor.style.foreach(_.dispose())
+    modelExecutor.executor.dispose()
     modelExecutor = null
 
     val (styleLoss, contentLoss) = getLoss(gram, content)

scala-package/examples/src/main/scala/org/apache/mxnetexamples/neuralstyle/end2end/GenV3.scala

@@ -22,18 +22,20 @@ import org.apache.mxnet.{Context, Shape, Symbol, Xavier}
 
 object GenV3 {
   def Conv(data: Symbol, numFilter: Int, kernel: (Int, Int) = (5, 5),
-      pad: (Int, Int) = (2, 2), stride: (Int, Int) = (2, 2)): Symbol = {
-    var sym = Symbol.api.Convolution(data = Some(data), num_filter = numFilter,
+           pad: (Int, Int) = (2, 2), stride: (Int, Int) = (2, 2)): Symbol = {
+    val sym1 = Symbol.api.Convolution(data = Some(data), num_filter = numFilter,
       kernel = Shape(kernel._1, kernel._2), stride = Some(Shape(stride._1, stride._2)),
       pad = Some(Shape(pad._1, pad._2)), no_bias = Some(false))
-    sym = Symbol.api.BatchNorm(data = Some(sym), fix_gamma = Some(false))
-    sym = Symbol.api.LeakyReLU(data = Some(sym), act_type = Some("leaky"))
-    sym
+    val sym2 = Symbol.api.BatchNorm(data = Some(sym1), fix_gamma = Some(false))
+    val sym3 = Symbol.api.LeakyReLU(data = Some(sym2), act_type = Some("leaky"))
+    sym2.dispose()
+    sym1.dispose()
+    sym3
   }
 
   def Deconv(data: Symbol, numFilter: Int, imHw: (Int, Int),
-      kernel: (Int, Int) = (7, 7), pad: (Int, Int) = (2, 2), stride: (Int, Int) = (2, 2),
-      crop: Boolean = true, out: Boolean = false): Symbol = {
+             kernel: (Int, Int) = (7, 7), pad: (Int, Int) = (2, 2), stride: (Int, Int) = (2, 2),
+             crop: Boolean = true, out: Boolean = false): Symbol = {
     var sym = Symbol.api.Deconvolution(data = Some(data), num_filter = numFilter,
       kernel = Shape(kernel._1, kernel._2), stride = Some(Shape(stride._1, stride._2)),
       pad = Some(Shape(pad._1, pad._2)), no_bias = Some(true))
@@ -74,9 +76,9 @@ object GenV3 {
       else (dataShape, false, false)
     }
     val mod = new Module(symbol = sym, context = ctx,
-                         dataShapes = dataShapes,
-                         initializer = new Xavier(magnitude = 2f),
-                         forTraining = forTraining, inputsNeedGrad = inputsNeedGrad)
+      dataShapes = dataShapes,
+      initializer = new Xavier(magnitude = 2f),
+      forTraining = forTraining, inputsNeedGrad = inputsNeedGrad)
     mod
   }
 }

scala-package/examples/src/main/scala/org/apache/mxnetexamples/neuralstyle/end2end/GenV4.scala

@@ -23,13 +23,15 @@ import org.apache.mxnet.{Context, Shape, Symbol, Xavier}
 object GenV4 {
 
   def Conv(data: Symbol, numFilter: Int, workspace : Long, kernel: (Int, Int) = (5, 5),
-      pad: (Int, Int) = (2, 2)): Symbol = {
-    var sym = Symbol.api.Convolution(data = Some(data), num_filter = numFilter,
+           pad: (Int, Int) = (2, 2)): Symbol = {
+    val sym1 = Symbol.api.Convolution(data = Some(data), num_filter = numFilter,
       kernel = Shape(kernel._1, kernel._2), workspace = Some(workspace),
       pad = Some(Shape(pad._1, pad._2)), no_bias = Some(false))
-    sym = Symbol.api.BatchNorm(data = Some(sym), fix_gamma = Some(false))
-    sym = Symbol.api.LeakyReLU(data = Some(sym), act_type = Some("leaky"))
-    sym
+    val sym2 = Symbol.api.BatchNorm(data = Some(sym1), fix_gamma = Some(false))
+    val sym3 = Symbol.api.LeakyReLU(data = Some(sym2), act_type = Some("leaky"))
+    sym2.dispose()
+    sym1.dispose()
+    sym3
   }
 
   def getGenerator(prefix: String, imHw: (Int, Int)): Symbol = {
@@ -41,7 +43,7 @@ object GenV4 {
     var conv4_1 = Conv(conv3_1, 32, 4096)
     var conv5_1 = Conv(conv4_1, 48, 4096)
     var conv6_1 = Conv(conv5_1, 32, 4096)
-    var out = Symbol.api.Convolution(data = Some(conv6_1), num_filter = 3, kernel = Shape(3, 3),
+    var out = Symbol.api.Convolution(data = Some(conv6_1), num_filter = 3, kernel = Shape(3,3),
       pad = Some(Shape(1, 1)), no_bias = Some(true), workspace = Some(4096))
     out = Symbol.api.BatchNorm(data = Some(out), fix_gamma = Some(false))
     out = Symbol.api.Activation(data = Some(out), act_type = "tanh")
@@ -62,9 +64,9 @@ object GenV4 {
       else (dataShape, false, false)
     }
     val mod = new Module(symbol = sym, context = ctx,
-                         dataShapes = dataShapes,
-                         initializer = new Xavier(magnitude = 2f),
-                         forTraining = forTraining, inputsNeedGrad = inputsNeedGrad)
+      dataShapes = dataShapes,
+      initializer = new Xavier(magnitude = 2f),
+      forTraining = forTraining, inputsNeedGrad = inputsNeedGrad)
     mod
   }
-}
+}