Repository: intel/ipex-llm
Pull request: fix: memory leak in `model.predictImageSet`.
## What changes were proposed in this pull request?
There're three reasons of memory leak.

1. repeat allocations in bigquant, which will be fixed in BigDL-core.
2. repeat clone module but no release. `model.predictImageSet` will new
   Predictor again and again.
3. share weights.

Because of multi models architecture, every `LocalPredictor` will clone models.
For efficiency all cloned models and origin model will share the same parameters.
If we iterate and release all cloned models, the parameters of origin model will
be released also. So we should clone a new model first in `LocalPredictor`.

This patch add a `StorageManager` which contains a concurrent hash map
to maintain all allocations of native memory/resources and prevent
duplicate release. It's also helpful for debug.

## How was this patch tested?

Add a new test case and run Jenkins.

## Related links or issues (optional)
fixed #2535



Problem statement:
Memory leak when using quantized model for prediction
The memory consumption grows steadily as multiple inferences running with quantized model. 

After some initial investigation, found the heap memory was actually small and consistent while the used physical memory is steadily increased.

Hints from issue discussion:
Please provide an example to reproduce the issue.
Synced with @i8run offline and he has successfully repro this issue.
User can reproduce it by calling `model.predictImage` multi times. Because of `predictImage` in `AbstractModule` will new `LocalPredictor` every call, which will `cloneModule` multi times. For normal model, the models cloned will be released by GC. But for quantized model, we should released it by manual.

Commit message: fix: move delete operation to ModelBroadcast

Patch to review:
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/models/utils/ModelBroadcast.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/models/utils/ModelBroadcast.scala
index cbda51eca05..153f1d72c5e 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/models/utils/ModelBroadcast.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/models/utils/ModelBroadcast.scala
@@ -16,14 +16,19 @@
 
 package com.intel.analytics.bigdl.models.utils
 
+import java.io.{IOException, ObjectInputStream, ObjectOutputStream}
+import java.util.UUID
+
 import com.intel.analytics.bigdl.Module
-import com.intel.analytics.bigdl.nn.{Container, Graph}
-import com.intel.analytics.bigdl.tensor.{QuantizedTensor, QuantizedType, Storage, Tensor}
+import com.intel.analytics.bigdl.nn.Container
+import com.intel.analytics.bigdl.nn.quantized.StorageManager
 import com.intel.analytics.bigdl.tensor.TensorNumericMath.TensorNumeric
+import com.intel.analytics.bigdl.tensor._
+import com.intel.analytics.bigdl.utils.Util._
 import org.apache.spark.SparkContext
 import org.apache.spark.broadcast.Broadcast
-import com.intel.analytics.bigdl.utils.Util._
 
+import scala.collection.mutable.ArrayBuffer
 import scala.reflect.ClassTag
 
 /**
@@ -38,10 +43,11 @@ import scala.reflect.ClassTag
 class ModelBroadcast[T: ClassTag](applyProtoBuffer: Boolean = false)
   (implicit ev: TensorNumeric[T]) extends Serializable {
 
-  private var broadcastModel: Broadcast[Module[T]] = _
+  private var broadcastModel: Broadcast[ModelInfo[T]] = _
   private var broadcastConsts: Broadcast[Map[String, Tensor[_]]] = _
   private var broadcastParameters: Broadcast[Array[Tensor[T]]] = _
 
+  val uuid: String = UUID.randomUUID().toString
 
   /**
    * broadcast the model
@@ -53,26 +59,45 @@ class ModelBroadcast[T: ClassTag](applyProtoBuffer: Boolean = false)
    * @return this
    */
   def broadcast(sc: SparkContext, model: Module[T]): this.type = {
+    CachedModels.deleteAll(uuid) // delete the models on driver
+
     if (applyProtoBuffer) {
-      broadcastModel = sc.broadcast(model)
+      broadcastModel = sc.broadcast(ModelInfo(uuid, model))
     } else {
+      // We should clone a new model which will maintain the origin model.
+      // Otherwise, the origin model's resources will be cleaned.
+      val newModel = model.cloneModule()
+      CachedModels.add(uuid, newModel)
+
       // broadcast Consts
-      if (model.isInstanceOf[Container[_, _, T]]) {
-        val moduleConsts = getAndClearConsts(model.asInstanceOf[Container[_, _, T]])
+      if (newModel.isInstanceOf[Container[_, _, T]]) {
+        val moduleConsts = getAndClearConsts(newModel.asInstanceOf[Container[_, _, T]])
         // TODO: broadcast Const, model structure and weight in the same broadcast.
         broadcastConsts = sc.broadcast(moduleConsts)
       }
+
       // broadcast weight and model
-      val weightsBias = getAndClearWeightBias(model.parameters())
-      broadcastModel = sc.broadcast(model.cloneModule())
+      val weightsBias = getAndClearWeightBias(newModel.parameters())
+
+      // We broadcast weight and model separately because of the memory limit of serialization.
+      // And we should clone the model structure (without weight) first because of lazy evaluation
+      // of broadcast. As you see, we have to put weights back to the model after broadcast call.
+      // As a quantized model, it will create relevant memory after clone because of
+      // `QuantizedTensor`. So we should release it first.
+      val cloned = newModel.cloneModule()
+      cloned.release()
+      CachedModels.add(uuid, cloned)
+
+      broadcastModel = sc.broadcast(ModelInfo[T](uuid, cloned))
       broadcastParameters = sc.broadcast(weightsBias)
 
-      putWeightBias(weightsBias, model)
-      initGradWeightBias(weightsBias, model)
+      putWeightBias(weightsBias, newModel)
+      initGradWeightBias(weightsBias, newModel)
     }
     this
   }
 
+
   /**
    * get the broadcast model
    * put the weight and bias back to the model
@@ -81,14 +106,21 @@ class ModelBroadcast[T: ClassTag](applyProtoBuffer: Boolean = false)
    * @return model
    */
   def value(initGradient: Boolean = false): Module[T] = {
+    CachedModels.deleteAll(uuid)
     if (applyProtoBuffer) {
-      val localModel = broadcastModel.value.clone(false)
+      val localModel = broadcastModel.value.model.clone(false)
+      val uuid = broadcastModel.value.uuid
+      CachedModels.add(uuid, localModel)
+
       if (initGradient) {
         initGradWeightBias(getWeightBias(localModel.parameters()), localModel)
       }
       localModel
     } else {
-      val localModel = broadcastModel.value.cloneModule()
+      val localModel = broadcastModel.value.model.cloneModule()
+      val uuid = broadcastModel.value.uuid
+      CachedModels.add(uuid, localModel)
+
       // share weight
       putWeightBias(broadcastParameters.value, localModel)
       // share Consts
@@ -141,13 +173,92 @@ class ModelBroadcast[T: ClassTag](applyProtoBuffer: Boolean = false)
       Array()
     }
   }
-
 }
 
-
 object ModelBroadcast {
   def apply[@specialized(Float, Double) T: ClassTag](applyProtoBuffer: Boolean = false)
         (implicit ev: TensorNumeric[T]) : ModelBroadcast[T] = {
     new ModelBroadcast(applyProtoBuffer)
   }
 }
+
+class ModelInfo[T: ClassTag](val uuid: String, @transient var model: Module[T])(
+  implicit ev: TensorNumeric[T]) extends Serializable {
+  @throws(classOf[IOException])
+  private def writeObject(out: ObjectOutputStream): Unit = {
+    out.defaultWriteObject()
+    val cloned = model.cloneModule()
+    out.writeObject(cloned)
+    CachedModels.add(uuid, cloned)
+  }
+}
+
+object ModelInfo {
+  def apply[T: ClassTag](uuid: String, model: Module[T])(
+    implicit ev: TensorNumeric[T]): ModelInfo[T] = new ModelInfo[T](uuid, model)
+}
+
+object CachedModels {
+  import java.util.concurrent.ConcurrentHashMap
+
+  import scala.collection._
+  import scala.collection.convert.decorateAsScala._
+
+  val lock = new AnyRef
+
+  type FloatValues = ArrayBuffer[Module[Float]]
+  type DoubleValues = ArrayBuffer[Module[Double]]
+
+  private val cachedFloatModels: concurrent.Map[String, FloatValues] =
+    new ConcurrentHashMap[String, FloatValues]().asScala
+  private val cachedDoubleModels: concurrent.Map[String, DoubleValues] =
+    new ConcurrentHashMap[String, DoubleValues]().asScala
+
+  def add[T: ClassTag](uuid: String, model: Module[T])( implicit ev: TensorNumeric[T]): Unit =
+    lock.synchronized {
+      ev.getType() match {
+        case FloatType =>
+          val models = cachedFloatModels.get(uuid) match {
+            case Some(values) => values += model.asInstanceOf[Module[Float]]
+            case _ => ArrayBuffer(model.asInstanceOf[Module[Float]])
+          }
+          cachedFloatModels.put(uuid, models)
+        case DoubleType =>
+          val models = cachedDoubleModels.get(uuid) match {
+            case Some(values) => values += model.asInstanceOf[Module[Double]]
+            case _ => ArrayBuffer(model.asInstanceOf[Module[Double]])
+          }
+          cachedDoubleModels.put(uuid, models)
+        case _ => throw new UnsupportedOperationException(s"unsupported type")
+      }
+    }
+
+  def deleteAll[T: ClassTag](currentKey: String)(implicit ev: TensorNumeric[T]): Unit =
+    lock.synchronized {
+      ev.getType() match {
+        case FloatType =>
+          val keys = cachedFloatModels.keys
+          for (key <- keys) {
+            if (key != currentKey) {
+              val models = cachedFloatModels(key)
+              println(s"delete key = $key ${models.length}")
+              for (model <- models) { model.release()
+                println(StorageManager.get().count(!_._2.isFreed))
+              }
+              cachedFloatModels.remove(key)
+              println(StorageManager.get().count(!_._2.isFreed))
+            }
+          }
+        case DoubleType =>
+          val keys = cachedDoubleModels.keys
+          for (key <- keys) {
+            if (key != currentKey) {
+              val models = cachedDoubleModels(key)
+              for (model <- models) { model.release() }
+              cachedDoubleModels.remove(key)
+            }
+          }
+        case _ => throw new UnsupportedOperationException(s"unsupported type")
+      }
+    }
+}
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/Container.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/Container.scala
index e48144a53dd..88f6049500c 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/Container.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/Container.scala
@@ -227,4 +227,8 @@ abstract class Container[A <: Activity : ClassTag,
     super.checkDuplicate(record)
     if (!skipDuplicateCheck()) modules.foreach(_.checkDuplicate(record))
   }
+
+  override def release(): Unit = {
+    modules.foreach(_.release())
+  }
 }
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/abstractnn/AbstractModule.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/abstractnn/AbstractModule.scala
index 391953ef9ea..cbabbf35d72 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/abstractnn/AbstractModule.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/abstractnn/AbstractModule.scala
@@ -665,8 +665,11 @@ abstract class AbstractModule[A <: Activity: ClassTag, B <: Activity: ClassTag,
         Predictor(this, featurePaddingParam, batchPerPartition)
           .predictImage(distributedImageFrame, outputLayer, shareBuffer, predictKey)
       case localImageFrame: LocalImageFrame =>
-        LocalPredictor(this, featurePaddingParam, batchPerPartition)
-          .predictImage(localImageFrame, outputLayer, shareBuffer, predictKey)
+        val predictor = LocalPredictor(this, featurePaddingParam, batchPerPartition)
+        val imageFrame = predictor.predictImage(localImageFrame, outputLayer, shareBuffer,
+          predictKey)
+        predictor.shutdown()
+        imageFrame
     }
   }
 
@@ -1094,5 +1097,11 @@ abstract class AbstractModule[A <: Activity: ClassTag, B <: Activity: ClassTag,
    * @return
    */
   private[nn] def skipDuplicateCheck(): Boolean = false
+
+  /**
+   * if the model contains native resources such as aligned memory, we should release it by manual.
+   * JVM GC can't release them reliably.
+   */
+  def release(): Unit = {}
 }
 
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/Desc.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/Desc.scala
index dbb63f0e656..a804c911620 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/Desc.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/Desc.scala
@@ -147,6 +147,9 @@ object Desc {
         desc
     }
 
+    // add every native memory allocation.
+    StorageManager.add(desc, params.getType)
+
     desc
   }
 
@@ -178,3 +181,28 @@ object QuantParams {
   val THRESHOLD = 127.0f
 }
 
+case class StorageInfo(descType: DescType, isFreed: Boolean)
+
+object StorageManager {
+  import java.util.concurrent.ConcurrentHashMap
+  private val nativeStorages: ConcurrentHashMap[Long, StorageInfo] = new ConcurrentHashMap()
+
+  def isFreed(nativeStorage: Long): Boolean = {
+    nativeStorages.get(nativeStorage).isFreed
+  }
+
+  // atomically set the value
+  def checkAndSet(nativeStorage: Long): Boolean = {
+    val descType = nativeStorages.get(nativeStorage).descType
+    nativeStorages.replace(nativeStorage, StorageInfo(descType, false), StorageInfo(descType, true))
+  }
+
+  def get(): Map[Long, StorageInfo] = {
+    import scala.collection.JavaConverters._
+    nativeStorages.asScala.toMap
+  }
+
+  def add(key: Long, descType: DescType): Unit = {
+    nativeStorages.put(key, StorageInfo(descType, false))
+  }
+}
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/Linear.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/Linear.scala
index 8063337378d..d13339859a5 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/Linear.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/Linear.scala
@@ -141,7 +141,7 @@ private[bigdl] class Linear[T: ClassTag](
     s"quantized.${getPrintName()}($inputSize -> $outputSize)"
   }
 
-  def release(): Unit = {
+  override def release(): Unit = {
     weight.release()
     data.release()
   }
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/SpatialConvolution.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/SpatialConvolution.scala
index 465d727ca2d..dec1ef10c1c 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/SpatialConvolution.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/SpatialConvolution.scala
@@ -270,7 +270,7 @@ private[bigdl] class SpatialConvolution[T: ClassTag](
       s" $kernelH, $strideW, $strideH, $padW, $padH, $nGroup)"
   }
 
-  def release(): Unit = {
+  override def release(): Unit = {
     weight.foreach(_.asInstanceOf[QuantizedTensor[T]].release())
     data.release()
   }
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/LocalPredictor.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/LocalPredictor.scala
index d0511fbfb8b..95b2d56afec 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/LocalPredictor.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/LocalPredictor.scala
@@ -18,11 +18,13 @@ package com.intel.analytics.bigdl.optim
 
 import com.intel.analytics.bigdl._
 import com.intel.analytics.bigdl.dataset.{SampleToMiniBatch, _}
+import com.intel.analytics.bigdl.nn.Container
 import com.intel.analytics.bigdl.nn.abstractnn.Activity
+import com.intel.analytics.bigdl.nn.quantized.QuantizedModule
 import com.intel.analytics.bigdl.tensor.TensorNumericMath.TensorNumeric
-import com.intel.analytics.bigdl.utils.{Engine, MklBlas, Util}
-import com.intel.analytics.bigdl.utils.Util._
 import com.intel.analytics.bigdl.transform.vision.image.{ImageFeature, ImageFrame, LocalImageFrame}
+import com.intel.analytics.bigdl.utils.Util._
+import com.intel.analytics.bigdl.utils.{Engine, MklBlas, Util}
 import org.apache.log4j.Logger
 
 import scala.reflect.ClassTag
@@ -58,15 +60,19 @@ class LocalPredictor[T: ClassTag] private[optim](model: Module[T],
     case _ => throw new IllegalArgumentException
   }
 
+  // we should clone a new model which has no impact to origin model
+  private val clonedModel = model.cloneModule()
+
   private val workingModels = {
-    val weightsBias = Util.getAndClearWeightBias(model.parameters())
+
+    val weightsBias = Util.getAndClearWeightBias(clonedModel.parameters())
     val models = (1 to subModelNumber).map(_ => {
-      val submodel = model.cloneModule().evaluate()
+      val submodel = clonedModel.cloneModule().evaluate()
       putWeightBias(weightsBias, submodel)
       submodel
     }).toArray
-    Util.putWeightBias(weightsBias, model)
-    Util.initGradWeightBias(weightsBias, model)
+    Util.putWeightBias(weightsBias, clonedModel)
+    Util.initGradWeightBias(weightsBias, clonedModel)
     models
   }
 
@@ -176,6 +182,14 @@ class LocalPredictor[T: ClassTag] private[optim](model: Module[T],
 
     ImageFrame.array(result.toArray)
   }
+
+  /**
+   * `shutdown` will release all native resources.
+   */
+  def shutdown(): Unit = {
+    workingModels.foreach(_.release())
+    clonedModel.release()
+  }
 }
 
 
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/Predictor.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/Predictor.scala
index d3d3f14e604..b6f6d79284d 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/Predictor.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/optim/Predictor.scala
@@ -16,9 +16,11 @@
 
 package com.intel.analytics.bigdl.optim
 
+import java.util.UUID
+
 import com.intel.analytics.bigdl._
 import com.intel.analytics.bigdl.dataset.{MiniBatch, PaddingParam, Sample, SampleToMiniBatch, Transformer, Utils, DataSet => _}
-import com.intel.analytics.bigdl.models.utils.ModelBroadcast
+import com.intel.analytics.bigdl.models.utils.{CachedModels, ModelBroadcast, ModelInfo}
 import com.intel.analytics.bigdl.nn.abstractnn.Activity
 import com.intel.analytics.bigdl.tensor.Tensor
 import com.intel.analytics.bigdl.tensor.TensorNumericMath.TensorNumeric
@@ -161,6 +163,8 @@ class Predictor[T: ClassTag] private[optim](
       partitionNum = Some(partitionNum),
       featurePaddingParam = featurePaddingParam))
     dataSet.mapPartitions { partition =>
+      CachedModels.add(modelBroad.uuid, model)
+
       val localModel = modelBroad.value()
       val localTransformer = otherBroad.value.cloneTransformer()
       val miniBatch = localTransformer(partition)
@@ -192,6 +196,9 @@ class Predictor[T: ClassTag] private[optim](
       partitionNum = Some(partitionNum),
       featurePaddingParam = featurePaddingParam), shareBuffer)
     val result = rdd.mapPartitions(partition => {
+      // By default, the `model` will be deserialized on worker, which will create new resources.
+      CachedModels.add(modelBroad.uuid, model)
+
       val localModel = modelBroad.value()
       val localToBatch = toBatchBroad.value._1.cloneTransformer()
 
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/DenseTensor.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/DenseTensor.scala
index 9c42c065b7f..2537dc1381c 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/DenseTensor.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/DenseTensor.scala
@@ -2204,6 +2204,8 @@ private[tensor] class DenseTensor[@specialized T: ClassTag](
     this.apply1(a => ev.digamma(a))
   }
 
+  override def toQuantizedTensor: QuantizedTensor[T] =
+    throw new IllegalArgumentException("DenseTensor cannot be cast to QuantizedTensor")
 }
 
 object DenseTensor {
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/QuantizedTensor.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/QuantizedTensor.scala
index 3b1485511d5..5d8ed2bc0d3 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/QuantizedTensor.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/QuantizedTensor.scala
@@ -45,7 +45,7 @@ private[bigdl] class QuantizedTensor[T: ClassTag](
   }
 
   def release(): this.type = {
-    if (desc != 0) {
+    if (desc != 0 && StorageManager.checkAndSet(desc)) {
       BigQuant.FreeMemory(desc)
     }
     desc = 0L
@@ -270,6 +270,8 @@ private[bigdl] class QuantizedTensor[T: ClassTag](
 
   override def getTensorNumeric(): TensorNumeric[T] = ev
 
+  override def toQuantizedTensor: QuantizedTensor[T] = this.asInstanceOf[QuantizedTensor[T]]
+
   @throws(classOf[IOException])
   private def readObject(in: ObjectInputStream): Unit = {
     in.defaultReadObject()
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/SparseTensor.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/SparseTensor.scala
index 028cd545db5..73516fc8a1e 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/SparseTensor.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/SparseTensor.scala
@@ -1126,6 +1126,9 @@ private[tensor] class SparseTensor[@specialized(Float, Double) T: ClassTag](
 
   override def sumSquare(): T =
     throw new UnsupportedOperationException(s"SparseTensor: Unimplemented method")
+
+  override def toQuantizedTensor: QuantizedTensor[T] =
+    throw new IllegalArgumentException("SparseTensor cannot be cast to QuantizedTensor")
 }
 
 object SparseTensor{
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/Tensor.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/Tensor.scala
index 7bcb47122ed..545a53de393 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/Tensor.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/Tensor.scala
@@ -804,6 +804,8 @@ trait Tensor[T] extends Serializable with TensorMath[T] with Activity {
     }
     return false
   }
+
+  def toQuantizedTensor: QuantizedTensor[T]
 }
 
 /**
diff --git a/spark/dl/src/main/scala/com/intel/analytics/bigdl/utils/Util.scala b/spark/dl/src/main/scala/com/intel/analytics/bigdl/utils/Util.scala
index 81337e07c47..3b539b6fcd7 100644
--- a/spark/dl/src/main/scala/com/intel/analytics/bigdl/utils/Util.scala
+++ b/spark/dl/src/main/scala/com/intel/analytics/bigdl/utils/Util.scala
@@ -19,7 +19,7 @@ package com.intel.analytics.bigdl.utils
 import java.io._
 
 import com.intel.analytics.bigdl._
-import com.intel.analytics.bigdl.nn.{Container, Graph}
+import com.intel.analytics.bigdl.nn.Container
 import com.intel.analytics.bigdl.nn.tf.Const
 import com.intel.analytics.bigdl.tensor.TensorNumericMath.{NumericWildcard, TensorNumeric}
 import com.intel.analytics.bigdl.tensor._
@@ -166,6 +166,10 @@ object Util {
     var i = 0
     while (i < tensors.length) {
       if (tensors(i) != null) {
+        if (tensors(i).getTensorType == QuantizedType) {
+          tensors(i).toQuantizedTensor.release()
+        }
+
         tensors(i).set()
       }
       i += 1
@@ -179,10 +183,23 @@ object Util {
     var i = 0
     while (i < localWeightBias.length) {
       if (localWeightBias(i) != null) {
-        localWeightBias(i).set(broadcastWeightBias(i))
+        clearAndSet(localWeightBias(i), broadcastWeightBias(i))
       }
       i += 1
     }
+
+    def clearAndSet(old: Tensor[T], other: Tensor[T]): Unit = {
+      if (old.getTensorType == QuantizedType && other.getTensorType == QuantizedType) {
+        val quantOld = old.asInstanceOf[QuantizedTensor[T]]
+        val quantOther = other.asInstanceOf[QuantizedTensor[T]]
+
+        if (quantOld.getNativeStorage != quantOther.getNativeStorage) {
+          quantOld.release()
+        }
+      }
+
+      old.set(other)
+    }
   }
 
   private[bigdl] def initGradWeightBias[T: ClassTag](
diff --git a/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/LocalPredictorSpec.scala b/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/LocalPredictorSpec.scala
index f6d63111f9e..d82616df4a8 100644
--- a/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/LocalPredictorSpec.scala
+++ b/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/LocalPredictorSpec.scala
@@ -23,6 +23,7 @@ import com.intel.analytics.bigdl.dataset.{PaddingParam, Sample, SampleToMiniBatc
 import com.intel.analytics.bigdl.models.inception.Inception_v1_NoAuxClassifier
 import com.intel.analytics.bigdl.nn.abstractnn.Activity
 import com.intel.analytics.bigdl.nn._
+import com.intel.analytics.bigdl.nn.quantized.StorageManager
 import com.intel.analytics.bigdl.tensor.Tensor
 import com.intel.analytics.bigdl.transform.vision.image._
 import com.intel.analytics.bigdl.transform.vision.image.augmentation.{CenterCrop, ChannelNormalize, Resize}
@@ -303,4 +304,35 @@ class LocalPredictorSpec extends FlatSpec with Matchers with BeforeAndAfter {
       assert(imageFeatures(x - 1).predict().asInstanceOf[Table].length() == 2)
     })
   }
+
+  "local predictor shutdown" should "work properly" in {
+    import com.intel.analytics.bigdl.numeric.NumericFloat
+    val input = Tensor[Float](4, 3, 224, 224).rand(-1, 1)
+
+    val samples = (1 to 20).map(i => {
+      Sample(Tensor[Float](3, 224, 224).randn())
+    }).toArray
+    val imageFrame = ImageFrame.array((0 until 20).map(x => {
+      val im = ImageFeature()
+      im(ImageFeature.sample) = samples(x)
+      im
+    }).toArray)
+
+    val model = Inception_v1_NoAuxClassifier(1000)
+    val quant = model.quantize().evaluate()
+    val initNativeSize = StorageManager.get().count(x => !x._2.isFreed)
+
+    // has no memory issues
+    (0 to 4).foreach { _ =>
+      quant.predictImage(imageFrame).toLocal().array.map(_.predict().asInstanceOf[Tensor[Float]])
+      StorageManager.get().count(x => !x._2.isFreed) should be (initNativeSize)
+    }
+
+    // check the model can work again
+    quant.forward(input)
+    val quant2 = model.quantize().evaluate()
+    quant2.forward(input)
+
+    quant.output.toTensor[Float] should be (quant2.output.toTensor[Float])
+  }
 }
diff --git a/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/PredictorSpec.scala b/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/PredictorSpec.scala
index f8648550edf..f29fa9f8ff5 100644
--- a/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/PredictorSpec.scala
+++ b/spark/dl/src/test/scala/com/intel/analytics/bigdl/optim/PredictorSpec.scala
@@ -19,12 +19,15 @@ package com.intel.analytics.bigdl.optim
 import com.intel.analytics.bigdl.dataset.{PaddingParam, Sample}
 import com.intel.analytics.bigdl.models.inception.Inception_v1_NoAuxClassifier
 import com.intel.analytics.bigdl.models.lenet.LeNet5
+import com.intel.analytics.bigdl.models.utils.CachedModels
 import com.intel.analytics.bigdl.nn._
+import com.intel.analytics.bigdl.nn.quantized.{StorageInfo, StorageManager}
 import com.intel.analytics.bigdl.tensor.Tensor
 import com.intel.analytics.bigdl.transform.vision.image._
 import com.intel.analytics.bigdl.transform.vision.image.augmentation.{CenterCrop, ChannelNormalize, Resize}
-import com.intel.analytics.bigdl.utils.{Engine, Table}
+import com.intel.analytics.bigdl.utils.{Engine, LoggerFilter, Table}
 import com.intel.analytics.bigdl.utils.RandomGenerator._
+import org.apache.log4j.{Level, Logger}
 import org.apache.spark.{SparkConf, SparkContext}
 import org.scalatest.{BeforeAndAfter, FlatSpec, Matchers}
 
@@ -211,4 +214,31 @@ class PredictorSpec extends FlatSpec with Matchers with BeforeAndAfter{
       assert(imageFeatures(x - 1).predict().asInstanceOf[Table].length() == 2)
     })
   }
+
+  "model predict should have no memory leak" should "be correct" in {
+    import com.intel.analytics.bigdl.tensor.TensorNumericMath.TensorNumeric.NumericFloat
+    LoggerFilter.redirectSparkInfoLogs()
+    Logger.getLogger("com.intel.analytics.bigdl.optim").setLevel(Level.INFO)
+    RNG.setSeed(100)
+    val resource = getClass.getClassLoader.getResource("pascal/")
+    val imageFrame = ImageFrame.read(resource.getFile, sc) ->
+      Resize(256, 256) -> CenterCrop(224, 224) ->
+      ChannelNormalize(0.485f, 0.456f, 0.406f, 0.229f, 0.224f, 0.225f) ->
+      MatToTensor() -> ImageFrameToSample()
+    val model = Inception_v1_NoAuxClassifier(classNum = 1000)
+    val quant = model.quantize()
+    val init = StorageManager.get()
+    println(s"init count ${init.count(!_._2.isFreed)}")
+    var second: Map[Long, StorageInfo] = null
+    (0 until 20).foreach { i =>
+      val detection = quant.predictImage(imageFrame, batchPerPartition = 16).toDistributed()
+      detection.rdd.first()
+      detection.rdd.collect()
+      println("=" * 80)
+      println(StorageManager.get().count(!_._2.isFreed))
+      println("-" * 80)
+    }
+    CachedModels.deleteAll("")
+    StorageManager.get().count(!_._2.isFreed) should be (init.count(!_._2.isFreed))
+  }
 }


Code review comments:
diff_hunk: @@ -141,13 +173,92 @@ class ModelBroadcast[T: ClassTag](applyProtoBuffer: Boolean = false)
       Array()
     }
   }
-
 }
 
-
 object ModelBroadcast {
   def apply[@specialized(Float, Double) T: ClassTag](applyProtoBuffer: Boolean = false)
         (implicit ev: TensorNumeric[T]) : ModelBroadcast[T] = {
     new ModelBroadcast(applyProtoBuffer)
   }
 }
+
+class ModelInfo[T: ClassTag](val uuid: String, @transient var model: Module[T])(
+  implicit ev: TensorNumeric[T]) extends Serializable {
+  @throws(classOf[IOException])
+  private def writeObject(out: ObjectOutputStream): Unit = {
+    out.defaultWriteObject()
+    val cloned = model.cloneModule()
+    out.writeObject(cloned)
+    CachedModels.add(uuid, cloned)
+  }
+}
+
+object ModelInfo {
+  def apply[T: ClassTag](uuid: String, model: Module[T])(
+    implicit ev: TensorNumeric[T]): ModelInfo[T] = new ModelInfo[T](uuid, model)
+}
+
+object CachedModels {
+  import java.util.concurrent.ConcurrentHashMap
+
+  import scala.collection._
+  import scala.collection.convert.decorateAsScala._
+
+  val lock = new AnyRef
+
+  type FloatValues = ArrayBuffer[Module[Float]]
+  type DoubleValues = ArrayBuffer[Module[Double]]
+
+  private val cachedFloatModels: concurrent.Map[String, FloatValues] =
+    new ConcurrentHashMap[String, FloatValues]().asScala
+  private val cachedDoubleModels: concurrent.Map[String, DoubleValues] =
+    new ConcurrentHashMap[String, DoubleValues]().asScala
+
+  def add[T: ClassTag](uuid: String, model: Module[T])( implicit ev: TensorNumeric[T]): Unit =
+    lock.synchronized {
+      ev.getType() match {
+        case FloatType =>
+          val models = cachedFloatModels.get(uuid) match {
+            case Some(values) => values += model.asInstanceOf[Module[Float]]
+            case _ => ArrayBuffer(model.asInstanceOf[Module[Float]])
+          }
+          cachedFloatModels.put(uuid, models)
+        case DoubleType =>
+          val models = cachedDoubleModels.get(uuid) match {
+            case Some(values) => values += model.asInstanceOf[Module[Double]]
+            case _ => ArrayBuffer(model.asInstanceOf[Module[Double]])
+          }
+          cachedDoubleModels.put(uuid, models)
+        case _ => throw new UnsupportedOperationException(s"unsupported type")
+      }
+    }
+
+  def deleteAll[T: ClassTag](currentKey: String)(implicit ev: TensorNumeric[T]): Unit =
+    lock.synchronized {
+      ev.getType() match {
+        case FloatType =>
+          val keys = cachedFloatModels.keys
+          for (key <- keys) {
+            if (key != currentKey) {
+              val models = cachedFloatModels(key)
+              println(s"delete key = $key ${models.length}")
+              for (model <- models) { model.release()
+                println(StorageManager.get().count(!_._2.isFreed))
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/models/utils/ModelBroadcast.scala
text: @user1:
remove the `println` clause ?

diff_hunk: @@ -141,13 +173,92 @@ class ModelBroadcast[T: ClassTag](applyProtoBuffer: Boolean = false)
       Array()
     }
   }
-
 }
 
-
 object ModelBroadcast {
   def apply[@specialized(Float, Double) T: ClassTag](applyProtoBuffer: Boolean = false)
         (implicit ev: TensorNumeric[T]) : ModelBroadcast[T] = {
     new ModelBroadcast(applyProtoBuffer)
   }
 }
+
+class ModelInfo[T: ClassTag](val uuid: String, @transient var model: Module[T])(
+  implicit ev: TensorNumeric[T]) extends Serializable {
+  @throws(classOf[IOException])
+  private def writeObject(out: ObjectOutputStream): Unit = {
+    out.defaultWriteObject()
+    val cloned = model.cloneModule()
+    out.writeObject(cloned)
+    CachedModels.add(uuid, cloned)
+  }
+}
+
+object ModelInfo {
+  def apply[T: ClassTag](uuid: String, model: Module[T])(
+    implicit ev: TensorNumeric[T]): ModelInfo[T] = new ModelInfo[T](uuid, model)
+}
+
+object CachedModels {
+  import java.util.concurrent.ConcurrentHashMap
+
+  import scala.collection._
+  import scala.collection.convert.decorateAsScala._
+
+  val lock = new AnyRef
+
+  type FloatValues = ArrayBuffer[Module[Float]]
+  type DoubleValues = ArrayBuffer[Module[Double]]
+
+  private val cachedFloatModels: concurrent.Map[String, FloatValues] =
+    new ConcurrentHashMap[String, FloatValues]().asScala
+  private val cachedDoubleModels: concurrent.Map[String, DoubleValues] =
+    new ConcurrentHashMap[String, DoubleValues]().asScala
+
+  def add[T: ClassTag](uuid: String, model: Module[T])( implicit ev: TensorNumeric[T]): Unit =
+    lock.synchronized {
+      ev.getType() match {
+        case FloatType =>
+          val models = cachedFloatModels.get(uuid) match {
+            case Some(values) => values += model.asInstanceOf[Module[Float]]
+            case _ => ArrayBuffer(model.asInstanceOf[Module[Float]])
+          }
+          cachedFloatModels.put(uuid, models)
+        case DoubleType =>
+          val models = cachedDoubleModels.get(uuid) match {
+            case Some(values) => values += model.asInstanceOf[Module[Double]]
+            case _ => ArrayBuffer(model.asInstanceOf[Module[Double]])
+          }
+          cachedDoubleModels.put(uuid, models)
+        case _ => throw new UnsupportedOperationException(s"unsupported type")
+      }
+    }
+
+  def deleteAll[T: ClassTag](currentKey: String)(implicit ev: TensorNumeric[T]): Unit =
+    lock.synchronized {
+      ev.getType() match {
+        case FloatType =>
+          val keys = cachedFloatModels.keys
+          for (key <- keys) {
+            if (key != currentKey) {
+              val models = cachedFloatModels(key)
+              println(s"delete key = $key ${models.length}")
+              for (model <- models) { model.release()
+                println(StorageManager.get().count(!_._2.isFreed))
+              }
+              cachedFloatModels.remove(key)
+              println(StorageManager.get().count(!_._2.isFreed))
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/models/utils/ModelBroadcast.scala
text: @user1:
remove the `println` clause ?

diff_hunk: @@ -38,10 +43,11 @@ import scala.reflect.ClassTag
 class ModelBroadcast[T: ClassTag](applyProtoBuffer: Boolean = false)
   (implicit ev: TensorNumeric[T]) extends Serializable {
 
-  private var broadcastModel: Broadcast[Module[T]] = _
+  private var broadcastModel: Broadcast[ModelInfo[T]] = _
   private var broadcastConsts: Broadcast[Map[String, Tensor[_]]] = _
   private var broadcastParameters: Broadcast[Array[Tensor[T]]] = _
 
+  val uuid: String = UUID.randomUUID().toString
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/models/utils/ModelBroadcast.scala
text: @user1:
private?

@user1:
why it's a `val`, for different predictions, should we use different uuid?

diff_hunk: @@ -178,3 +181,28 @@ object QuantParams {
   val THRESHOLD = 127.0f
 }
 
+case class StorageInfo(descType: DescType, isFreed: Boolean)
+
+object StorageManager {
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/Desc.scala
text: @user1:
private[bigdl]?

diff_hunk: @@ -178,3 +181,28 @@ object QuantParams {
   val THRESHOLD = 127.0f
 }
 
+case class StorageInfo(descType: DescType, isFreed: Boolean)
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/nn/quantized/Desc.scala
text: @user1:
private[bigdl]?

diff_hunk: @@ -804,6 +804,8 @@ trait Tensor[T] extends Serializable with TensorMath[T] with Activity {
     }
     return false
   }
+
+  def toQuantizedTensor: QuantizedTensor[T]
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/tensor/Tensor.scala
text: @user1:
private[bigdl]?

diff_hunk: @@ -141,13 +173,92 @@ class ModelBroadcast[T: ClassTag](applyProtoBuffer: Boolean = false)
       Array()
     }
   }
-
 }
 
-
 object ModelBroadcast {
   def apply[@specialized(Float, Double) T: ClassTag](applyProtoBuffer: Boolean = false)
         (implicit ev: TensorNumeric[T]) : ModelBroadcast[T] = {
     new ModelBroadcast(applyProtoBuffer)
   }
 }
+
+class ModelInfo[T: ClassTag](val uuid: String, @transient var model: Module[T])(
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/models/utils/ModelBroadcast.scala
text: @user1:
case class ?

@user1:
Why transient ? will it be ignored while serialized and deserialized between driver and executor ?

@user2:
private[bigdl]?

diff_hunk: @@ -141,13 +173,92 @@ class ModelBroadcast[T: ClassTag](applyProtoBuffer: Boolean = false)
       Array()
     }
   }
-
 }
 
-
 object ModelBroadcast {
   def apply[@specialized(Float, Double) T: ClassTag](applyProtoBuffer: Boolean = false)
         (implicit ev: TensorNumeric[T]) : ModelBroadcast[T] = {
     new ModelBroadcast(applyProtoBuffer)
   }
 }
+
+class ModelInfo[T: ClassTag](val uuid: String, @transient var model: Module[T])(
+  implicit ev: TensorNumeric[T]) extends Serializable {
+  @throws(classOf[IOException])
+  private def writeObject(out: ObjectOutputStream): Unit = {
+    out.defaultWriteObject()
+    val cloned = model.cloneModule()
+    out.writeObject(cloned)
+    CachedModels.add(uuid, cloned)
+  }
+}
+
+object ModelInfo {
+  def apply[T: ClassTag](uuid: String, model: Module[T])(
+    implicit ev: TensorNumeric[T]): ModelInfo[T] = new ModelInfo[T](uuid, model)
+}
+
+object CachedModels {
+  import java.util.concurrent.ConcurrentHashMap
+
+  import scala.collection._
+  import scala.collection.convert.decorateAsScala._
+
+  val lock = new AnyRef
+
+  type FloatValues = ArrayBuffer[Module[Float]]
+  type DoubleValues = ArrayBuffer[Module[Double]]
+
+  private val cachedFloatModels: concurrent.Map[String, FloatValues] =
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/models/utils/ModelBroadcast.scala
text: @user1:
Why there're two arraybuffer here?

diff_hunk: @@ -141,13 +173,92 @@ class ModelBroadcast[T: ClassTag](applyProtoBuffer: Boolean = false)
       Array()
     }
   }
-
 }
 
-
 object ModelBroadcast {
   def apply[@specialized(Float, Double) T: ClassTag](applyProtoBuffer: Boolean = false)
         (implicit ev: TensorNumeric[T]) : ModelBroadcast[T] = {
     new ModelBroadcast(applyProtoBuffer)
   }
 }
+
+class ModelInfo[T: ClassTag](val uuid: String, @transient var model: Module[T])(
+  implicit ev: TensorNumeric[T]) extends Serializable {
+  @throws(classOf[IOException])
+  private def writeObject(out: ObjectOutputStream): Unit = {
+    out.defaultWriteObject()
+    val cloned = model.cloneModule()
+    out.writeObject(cloned)
+    CachedModels.add(uuid, cloned)
+  }
+}
+
+object ModelInfo {
+  def apply[T: ClassTag](uuid: String, model: Module[T])(
+    implicit ev: TensorNumeric[T]): ModelInfo[T] = new ModelInfo[T](uuid, model)
+}
+
+object CachedModels {
+  import java.util.concurrent.ConcurrentHashMap
+
+  import scala.collection._
+  import scala.collection.convert.decorateAsScala._
+
+  val lock = new AnyRef
+
+  type FloatValues = ArrayBuffer[Module[Float]]
+  type DoubleValues = ArrayBuffer[Module[Double]]
+
+  private val cachedFloatModels: concurrent.Map[String, FloatValues] =
+    new ConcurrentHashMap[String, FloatValues]().asScala
+  private val cachedDoubleModels: concurrent.Map[String, DoubleValues] =
+    new ConcurrentHashMap[String, DoubleValues]().asScala
+
+  def add[T: ClassTag](uuid: String, model: Module[T])( implicit ev: TensorNumeric[T]): Unit =
+    lock.synchronized {
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/models/utils/ModelBroadcast.scala
text: @user1:
CachedModels.sync?

diff_hunk: @@ -141,13 +173,92 @@ class ModelBroadcast[T: ClassTag](applyProtoBuffer: Boolean = false)
       Array()
     }
   }
-
 }
 
-
 object ModelBroadcast {
   def apply[@specialized(Float, Double) T: ClassTag](applyProtoBuffer: Boolean = false)
         (implicit ev: TensorNumeric[T]) : ModelBroadcast[T] = {
     new ModelBroadcast(applyProtoBuffer)
   }
 }
+
+class ModelInfo[T: ClassTag](val uuid: String, @transient var model: Module[T])(
+  implicit ev: TensorNumeric[T]) extends Serializable {
+  @throws(classOf[IOException])
+  private def writeObject(out: ObjectOutputStream): Unit = {
+    out.defaultWriteObject()
+    val cloned = model.cloneModule()
+    out.writeObject(cloned)
+    CachedModels.add(uuid, cloned)
+  }
+}
+
+object ModelInfo {
+  def apply[T: ClassTag](uuid: String, model: Module[T])(
+    implicit ev: TensorNumeric[T]): ModelInfo[T] = new ModelInfo[T](uuid, model)
+}
+
+object CachedModels {
+  import java.util.concurrent.ConcurrentHashMap
+
+  import scala.collection._
+  import scala.collection.convert.decorateAsScala._
+
+  val lock = new AnyRef
+
+  type FloatValues = ArrayBuffer[Module[Float]]
+  type DoubleValues = ArrayBuffer[Module[Double]]
+
+  private val cachedFloatModels: concurrent.Map[String, FloatValues] =
+    new ConcurrentHashMap[String, FloatValues]().asScala
+  private val cachedDoubleModels: concurrent.Map[String, DoubleValues] =
+    new ConcurrentHashMap[String, DoubleValues]().asScala
+
+  def add[T: ClassTag](uuid: String, model: Module[T])( implicit ev: TensorNumeric[T]): Unit =
+    lock.synchronized {
+      ev.getType() match {
+        case FloatType =>
+          val models = cachedFloatModels.get(uuid) match {
+            case Some(values) => values += model.asInstanceOf[Module[Float]]
+            case _ => ArrayBuffer(model.asInstanceOf[Module[Float]])
+          }
+          cachedFloatModels.put(uuid, models)
+        case DoubleType =>
+          val models = cachedDoubleModels.get(uuid) match {
+            case Some(values) => values += model.asInstanceOf[Module[Double]]
+            case _ => ArrayBuffer(model.asInstanceOf[Module[Double]])
+          }
+          cachedDoubleModels.put(uuid, models)
+        case _ => throw new UnsupportedOperationException(s"unsupported type")
+      }
+    }
+
+  def deleteAll[T: ClassTag](currentKey: String)(implicit ev: TensorNumeric[T]): Unit =
+    lock.synchronized {
+      ev.getType() match {
+        case FloatType =>
+          val keys = cachedFloatModels.keys
+          for (key <- keys) {
+            if (key != currentKey) {
+              val models = cachedFloatModels(key)
+              println(s"delete key = $key ${models.length}")
+              for (model <- models) { model.release()
path: spark/dl/src/main/scala/com/intel/analytics/bigdl/models/utils/ModelBroadcast.scala
text: @user1:
Does `release` compatible for those models not quantized ?