Skip to content

Commit

Permalink
fix: memory leak in model.predictImageSet. (intel-analytics#2557)
Browse files Browse the repository at this point in the history
* fix: memory leak in `model.predictImageSet`.

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.
2. share weights.

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.

* fix: delete .

* refator:  as the API for AbstractModule

* fix: distribute predictor memory leak

* fix: move delete operation to ModelBroadcast

* refinement per review

* fix ut

* fix scala version issue
  • Loading branch information
i8run authored and wzhongyuan committed Jun 28, 2018
1 parent 229c5df commit 2f6b5e6
Showing 1 changed file with 19 additions and 2 deletions.
21 changes: 19 additions & 2 deletions dl/src/main/scala/com/intel/analytics/bigdl/common/Util.scala
Original file line number Diff line number Diff line change
Expand Up @@ -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._
Expand Down Expand Up @@ -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
Expand All @@ -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](
Expand Down

0 comments on commit 2f6b5e6

Please sign in to comment.