From 9b153413aa7dee24a0f70bdf13dc0380fb5e7d30 Mon Sep 17 00:00:00 2001 From: Yizhi Liu Date: Thu, 12 Jul 2018 17:24:43 -0700 Subject: [PATCH 1/9] [Scala] NDArrayCollector for automatically disposing NDArrays --- .../scala/org/apache/mxnet/Executor.scala | 2 +- .../main/scala/org/apache/mxnet/Monitor.scala | 2 +- .../main/scala/org/apache/mxnet/NDArray.scala | 13 +- .../org/apache/mxnet/NDArrayCollector.scala | 162 ++++++++++++++++++ .../scala/org/apache/mxnet/Operator.scala | 4 +- .../org/apache/mxnet/io/NDArrayIter.scala | 17 +- .../apache/mxnet/NDArrayCollectorSuite.scala | 67 ++++++++ 7 files changed, 252 insertions(+), 15 deletions(-) create mode 100644 scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala create mode 100644 scala-package/core/src/test/scala/org/apache/mxnet/NDArrayCollectorSuite.scala diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/Executor.scala b/scala-package/core/src/main/scala/org/apache/mxnet/Executor.scala index 2f79b58a52cd..181b2328ddcc 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/Executor.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/Executor.scala @@ -167,7 +167,7 @@ class Executor private[mxnet](private[mxnet] val handle: ExecutorHandle, private def getOutputs: Array[NDArray] = { val ndHandles = ArrayBuffer[NDArrayHandle]() checkCall(_LIB.mxExecutorOutputs(handle, ndHandles)) - ndHandles.toArray.map(new NDArray(_)) + ndHandles.toArray.map(new NDArray(_, addToCollector = false)) } /** diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/Monitor.scala b/scala-package/core/src/main/scala/org/apache/mxnet/Monitor.scala index 8e53d652fde6..c8a251d03a6c 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/Monitor.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/Monitor.scala @@ -51,7 +51,7 @@ class Monitor( override def invoke(name: String, arr: NDArrayHandle): Unit = { // wrapper for executor callback if (activated) { - val array = new NDArray(arr, writable = false) + val array = new NDArray(arr, writable = false, addToCollector = false) val elem = (step, name, statFunc(array)) queue += elem } diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala b/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala index c2de6ea43f2c..dac33118e1c1 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala @@ -100,6 +100,7 @@ object NDArray extends NDArrayBase { val outputs = ArrayBuffer.empty[NDArrayHandle] checkCall(_LIB.mxImperativeInvoke(function.handle, ndArgs.map(_.handle).toArray, outputVars, outputs, updatedKwargs.size, updatedKwargs.keys.toArray, updatedKwargs.values.toArray)) + new NDArrayFuncReturn(Option(oriOutputs).getOrElse { val outputArrs = outputs.map(new NDArray(_)).toArray addDependency(ndArgs.toArray, outputArrs) @@ -490,7 +491,7 @@ object NDArray extends NDArrayBase { val names = ArrayBuffer.empty[String] checkCall(_LIB.mxNDArrayLoad(fname, outSize, handles, outNameSize, names)) require(outNameSize.value == 0 || outNameSize.value == outSize.value) - (names.toArray, handles.map(new NDArray(_)).toArray) + (names.toArray, handles.map { ndHandle => new NDArray(ndHandle) }.toArray) } def load2Map(fname: String): Map[String, NDArray] = { @@ -554,11 +555,16 @@ object NDArray extends NDArrayBase { * */ class NDArray private[mxnet](private[mxnet] val handle: NDArrayHandle, - val writable: Boolean = true) extends WarnIfNotDisposed { + val writable: Boolean = true, + addToCollector: Boolean = true) extends WarnIfNotDisposed { + if (addToCollector) { + NDArrayCollector.collect(this) + } + // record arrays who construct this array instance // we use weak reference to prevent gc blocking private[mxnet] val dependencies = mutable.HashMap.empty[Long, WeakReference[NDArray]] - private var disposed = false + @volatile private var disposed = false def isDisposed: Boolean = disposed def serialize(): Array[Byte] = { @@ -979,6 +985,7 @@ class NDArray private[mxnet](private[mxnet] val handle: NDArrayHandle, def copyTo(ctx: Context): NDArray = { val ret = new NDArray(NDArray.newAllocHandle(shape, ctx, delayAlloc = true)) copyTo(ret) + ret } /** diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala new file mode 100644 index 000000000000..f05b1b24b398 --- /dev/null +++ b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala @@ -0,0 +1,162 @@ +/* + * 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 org.apache.mxnet + +import org.slf4j.LoggerFactory + +import scala.annotation.varargs +import scala.collection.mutable + +/** + * A collector to store NDArrays. + * It provides a scope, NDArrays allocated in the scope can either
+ * - be disposed automatically when the scope finished, or
+ * - simply be collected for future usage. + *
+ * NDArrayCollector.auto or NDArrayCollector.later is thread-safe, + * i.e., one can use them in different threads, + * the NDArrays will be stored in the collector of current thread, + * preventing from sharing across the threads. + *
+ * If the return type of scope is NDArray or NDArrayFuncReturn, + * the collector is smart enough NOT to collect or dispose the returned NDArray.
+ * However in other cases, it is users' responsibility NOT to leak allocated NDArrays outside, + * (e.g., store to a global variable and use later, pass to another thread, etc.)
+ * Usage Example: + *
+ *  val a = NDArray.array(Array(-1f, 0f, 1f, 2f, 3f, 4f), shape = Shape(2, 3))
+ *  val res = NDArrayCollector.auto().withScope {
+ *    (NDArray.relu(a) + a).toArray
+ *  }
+ *  
+ * In the case above, the intermediate NDArrays + * (created by NDArray.relu and +) will be disposed automatically.
+ * User can also decide to use dispose the collected NDArrays later:
+ *
+ *  val collector = NDArrayCollector.manual()
+ *  val res = collector.withScope {
+ *    (NDArray.relu(a) + a).toArray
+ *  }
+ *  collector.foreach(_.dispose())
+ *  
+ * For Java users:
+ *
+ *  NDArray a = NDArray.array(new float[]{-1f, 0f, 1f, 2f, 3f, 4f},
+ *                            Shape.create(2, 3), Context.cpu(0));
+ *  float[] sliced = NDArrayCollector.auto().withScope(
+ *    new scala.runtime.AbstractFunction0() {
+ *    @Override
+ *    public float[] apply() {
+ *      a.slice(0, 1).toArray();
+ *    }
+ *  });
+ *  
+ */ +object NDArrayCollector { + private val logger = LoggerFactory.getLogger(classOf[NDArrayCollector]) + + private val currCollector = new ThreadLocal[NDArrayCollector] { + override def initialValue = new NDArrayCollector(false, false) + } + + /** + * Create a collector which will dispose the collected NDArrays automatically. + * @return an auto-disposable collector. + */ + def auto(): NDArrayCollector = new NDArrayCollector(true) + + /** + * Create a collector allows users to later dispose the collected NDArray manually. + * @return a manually-disposable collector. + */ + def manual(): NDArrayCollector = new NDArrayCollector(false) + + /** + * Collect the NDArrays into the collector of the current thread. + * @param ndArray NDArrays need to be collected. + */ + @varargs def collect(ndArray: NDArray*): Unit = { + ndArray.foreach(nd => currCollector.get().add(nd)) + } +} + +class NDArrayCollector private(private val autoDispose: Boolean = true, + private val doCollect: Boolean = true) { + private val arrays: mutable.Map[Long, NDArray] = mutable.HashMap.empty[Long, NDArray] + + private def add(nd: NDArray*): Unit = { + if (doCollect) nd.foreach(arr => arrays.put(arr.handle, arr)) + } + + /** + * Clear the collector. + */ + def clear(): Unit = { + arrays.clear() + } + + /** + * Iterate over the collected NDArrays and apply the user-defined function to each NDArray. + * @param f the function that is applied for its side-effect to every NDArray. + * The result of function f is discarded. + */ + def foreach(f: NDArray => Unit): Unit = { + arrays.values.foreach(f(_)) + } + + /** + * @return how many unique NDArrays are collected. + */ + def size: Int = arrays.size + + /** + * Create a code scope, NDArrays allocated within this scope will be collected. + * The collected NDArrays will be either
+ * - disposed automatically when the scope finishes (when using auto) or
+ * - stored for later access (when using manual)
+ * If the return type of scope is NDArray or NDArrayFuncReturn, + * it is smart enough NOT to collect or dispose the returned NDArray.
+ * However in other cases, it is users' responsibility NOT to leak allocated NDArrays outside. + * @param body function to be executed within the scope. + * @tparam T return type of the function body. + * @return The result of function body. + */ + def withScope[T](body: => T): T = { + val old = NDArrayCollector.currCollector.get() + NDArrayCollector.currCollector.set(this) + try { + val ret = body + + ret match { + case ndRet: NDArray => + arrays.remove(ndRet.handle) + case ndarrays: NDArrayFuncReturn => + ndarrays.arr.foreach(nd => arrays.remove(nd.handle)) + case _ => // do nothing + } + + if (autoDispose) { + foreach(_.dispose()) + clear() + } + ret + } finally { + NDArrayCollector.currCollector.set(old) + } + } +} diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/Operator.scala b/scala-package/core/src/main/scala/org/apache/mxnet/Operator.scala index 6630d5ff53d3..f2abe5e45159 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/Operator.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/Operator.scala @@ -72,9 +72,9 @@ abstract class CustomOp { val tensors = (0 until 5).toArray.map( x => ArrayBuffer[NDArray]() ) for (i <- 0 until numNdarray) { if (tags(i) == 1 || tags(i) == 4) { - tensors(tags(i)) += new NDArray(ndarraies(i), writable = true) + tensors(tags(i)) += new NDArray(ndarraies(i), writable = true, addToCollector = false) } else { - tensors(tags(i)) += new NDArray(ndarraies(i), writable = false) + tensors(tags(i)) += new NDArray(ndarraies(i), writable = false, addToCollector = false) } } val reqEnum = Array("null", "write", "inplace", "add") diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/io/NDArrayIter.scala b/scala-package/core/src/main/scala/org/apache/mxnet/io/NDArrayIter.scala index 70c648778870..10461315c198 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/io/NDArrayIter.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/io/NDArrayIter.scala @@ -161,14 +161,15 @@ class NDArrayIter(data: IndexedSeq[(String, NDArray)], */ private def _padData(ndArray: NDArray): NDArray = { val padNum = cursor + dataBatchSize - numData - val newArray = NDArray.zeros(ndArray.slice(0, dataBatchSize).shape) - val batch = ndArray.slice(cursor, numData) - val padding = ndArray.slice(0, padNum) - newArray.slice(0, dataBatchSize - padNum).set(batch).dispose() - newArray.slice(dataBatchSize - padNum, dataBatchSize).set(padding).dispose() - batch.dispose() - padding.dispose() - newArray + val shape = Shape(dataBatchSize) ++ ndArray.shape.slice(1, ndArray.shape.size) + val newArray = NDArray.zeros(shape) + NDArrayCollector.auto().withScope { + val batch = ndArray.slice(cursor, numData) + val padding = ndArray.slice(0, padNum) + newArray.slice(0, dataBatchSize - padNum).set(batch) + newArray.slice(dataBatchSize - padNum, dataBatchSize).set(padding) + newArray + } } private def _getData(data: IndexedSeq[(String, NDArray)]): IndexedSeq[NDArray] = { diff --git a/scala-package/core/src/test/scala/org/apache/mxnet/NDArrayCollectorSuite.scala b/scala-package/core/src/test/scala/org/apache/mxnet/NDArrayCollectorSuite.scala new file mode 100644 index 000000000000..55fa26b2b0c1 --- /dev/null +++ b/scala-package/core/src/test/scala/org/apache/mxnet/NDArrayCollectorSuite.scala @@ -0,0 +1,67 @@ +/* + * 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 org.apache.mxnet + +import org.scalatest.{BeforeAndAfterAll, FunSuite, Matchers} + +class NDArrayCollectorSuite extends FunSuite with BeforeAndAfterAll with Matchers { + + test("auto dispose") { + val a = NDArray.array(Array(-1f, 0f, 1f, 2f, 3f, 4f), shape = Shape(2, 3)) + var b, c: NDArray = null + + val res = NDArrayCollector.auto().withScope { + b = NDArray.relu(a) // [0, 0, 1, 2, 3, 4] + c = a + b // [-1, 0, 2, 4, 6, 8] + c.slice(0, 1) + } + + assert(b.isDisposed) + assert(c.isDisposed) + assert(!res.isDisposed) // smart enough not to dispose the returned NDArray + + assert(res.toArray === Array(-1f, 0f, 2f)) + } + + test("manually dispose") { + val a = NDArray.array(Array(-1f, 0f, 1f, 2f, 3f, 4f), shape = Shape(2, 3)) + var b, c: NDArray = null + + val collector = NDArrayCollector.manual() + val res = collector.withScope { + b = NDArray.relu(a) // [0, 0, 1, 2, 3, 4] + c = a + b // [-1, 0, 2, 4, 6, 8] + c.slice(0, 1) + } + + assert(res.toArray === Array(-1f, 0f, 2f)) + + assert(collector.size === 2) // smart enough not to collect the returned NDArray + assert(!b.isDisposed) + assert(!c.isDisposed) + assert(!res.isDisposed) + + collector.foreach(_.dispose()) + assert(b.isDisposed) + assert(c.isDisposed) + assert(!res.isDisposed) + + collector.clear() + assert(collector.size === 0) + } +} From 6ab02a3de87ca7ba7ea53cd8f2498f805628e8e2 Mon Sep 17 00:00:00 2001 From: Yizhi Liu Date: Fri, 13 Jul 2018 10:33:18 -0700 Subject: [PATCH 2/9] modify doc for NDArrayCollector --- .../src/main/scala/org/apache/mxnet/NDArrayCollector.scala | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala index f05b1b24b398..cf8cccd2a8cc 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala @@ -28,11 +28,6 @@ import scala.collection.mutable * - be disposed automatically when the scope finished, or
* - simply be collected for future usage. *
- * NDArrayCollector.auto or NDArrayCollector.later is thread-safe, - * i.e., one can use them in different threads, - * the NDArrays will be stored in the collector of current thread, - * preventing from sharing across the threads. - *
* If the return type of scope is NDArray or NDArrayFuncReturn, * the collector is smart enough NOT to collect or dispose the returned NDArray.
* However in other cases, it is users' responsibility NOT to leak allocated NDArrays outside, @@ -91,7 +86,7 @@ object NDArrayCollector { * @param ndArray NDArrays need to be collected. */ @varargs def collect(ndArray: NDArray*): Unit = { - ndArray.foreach(nd => currCollector.get().add(nd)) + currCollector.get().add(ndArray: _*) } } From 8addba80e639e5a23da0a1efafcb93d8468f7657 Mon Sep 17 00:00:00 2001 From: Yizhi Liu Date: Fri, 13 Jul 2018 10:47:38 -0700 Subject: [PATCH 3/9] modify the function doc of NDArrayCollector.withScope --- .../src/main/scala/org/apache/mxnet/NDArrayCollector.scala | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala index cf8cccd2a8cc..3522b50092bb 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala @@ -25,7 +25,7 @@ import scala.collection.mutable /** * A collector to store NDArrays. * It provides a scope, NDArrays allocated in the scope can either
- * - be disposed automatically when the scope finished, or
+ * - be disposed automatically when the code block finishes, or
* - simply be collected for future usage. *
* If the return type of scope is NDArray or NDArrayFuncReturn, @@ -122,12 +122,12 @@ class NDArrayCollector private(private val autoDispose: Boolean = true, /** * Create a code scope, NDArrays allocated within this scope will be collected. * The collected NDArrays will be either
- * - disposed automatically when the scope finishes (when using auto) or
+ * - disposed automatically when the code blcok finishes (when using auto) or
* - stored for later access (when using manual)
* If the return type of scope is NDArray or NDArrayFuncReturn, * it is smart enough NOT to collect or dispose the returned NDArray.
* However in other cases, it is users' responsibility NOT to leak allocated NDArrays outside. - * @param body function to be executed within the scope. + * @param body code block to be executed within the scope. * @tparam T return type of the function body. * @return The result of function body. */ From 731066dcf906a4953ae05ac2d8d5e07cdd63e1d4 Mon Sep 17 00:00:00 2001 From: Yizhi Liu Date: Fri, 13 Jul 2018 11:21:16 -0700 Subject: [PATCH 4/9] remove trivial changes --- .../core/src/main/scala/org/apache/mxnet/NDArray.scala | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala b/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala index dac33118e1c1..58ab5cadd9d5 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/NDArray.scala @@ -100,7 +100,6 @@ object NDArray extends NDArrayBase { val outputs = ArrayBuffer.empty[NDArrayHandle] checkCall(_LIB.mxImperativeInvoke(function.handle, ndArgs.map(_.handle).toArray, outputVars, outputs, updatedKwargs.size, updatedKwargs.keys.toArray, updatedKwargs.values.toArray)) - new NDArrayFuncReturn(Option(oriOutputs).getOrElse { val outputArrs = outputs.map(new NDArray(_)).toArray addDependency(ndArgs.toArray, outputArrs) @@ -491,7 +490,7 @@ object NDArray extends NDArrayBase { val names = ArrayBuffer.empty[String] checkCall(_LIB.mxNDArrayLoad(fname, outSize, handles, outNameSize, names)) require(outNameSize.value == 0 || outNameSize.value == outSize.value) - (names.toArray, handles.map { ndHandle => new NDArray(ndHandle) }.toArray) + (names.toArray, handles.map(new NDArray(_)).toArray) } def load2Map(fname: String): Map[String, NDArray] = { @@ -985,7 +984,6 @@ class NDArray private[mxnet](private[mxnet] val handle: NDArrayHandle, def copyTo(ctx: Context): NDArray = { val ret = new NDArray(NDArray.newAllocHandle(shape, ctx, delayAlloc = true)) copyTo(ret) - ret } /** From ddf5daafd8d67b246aa97cf9c82e42ee26688940 Mon Sep 17 00:00:00 2001 From: Yizhi Liu Date: Fri, 13 Jul 2018 12:55:36 -0700 Subject: [PATCH 5/9] put dispose in finally --- .../src/main/scala/org/apache/mxnet/NDArrayCollector.scala | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala index 3522b50092bb..bbbca23d4174 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala @@ -136,7 +136,6 @@ class NDArrayCollector private(private val autoDispose: Boolean = true, NDArrayCollector.currCollector.set(this) try { val ret = body - ret match { case ndRet: NDArray => arrays.remove(ndRet.handle) @@ -144,13 +143,12 @@ class NDArrayCollector private(private val autoDispose: Boolean = true, ndarrays.arr.foreach(nd => arrays.remove(nd.handle)) case _ => // do nothing } - + ret + } finally { if (autoDispose) { foreach(_.dispose()) clear() } - ret - } finally { NDArrayCollector.currCollector.set(old) } } From 72cf085e4b78c9be4087f983b7ed7788e628da13 Mon Sep 17 00:00:00 2001 From: Yizhi Liu Date: Fri, 13 Jul 2018 16:09:59 -0700 Subject: [PATCH 6/9] fix jni NDArray signature --- .../native/src/main/native/org_apache_mxnet_native_c_api.cc | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/scala-package/native/src/main/native/org_apache_mxnet_native_c_api.cc b/scala-package/native/src/main/native/org_apache_mxnet_native_c_api.cc index caf7e56cd3bd..d944a8d049ca 100644 --- a/scala-package/native/src/main/native/org_apache_mxnet_native_c_api.cc +++ b/scala-package/native/src/main/native/org_apache_mxnet_native_c_api.cc @@ -612,7 +612,7 @@ extern "C" void KVStoreUpdaterCallbackFunc // find java NDArray constructor jclass ndObjClass = env->FindClass("org/apache/mxnet/NDArray"); - jmethodID ndObjConstructor = env->GetMethodID(ndObjClass, "", "(JZ)V"); + jmethodID ndObjConstructor = env->GetMethodID(ndObjClass, "", "(JZZ)V"); jobject ndRecv = env->NewObject(ndObjClass, ndObjConstructor, reinterpret_cast(recv), true); From 1417472f1eaa70424e609ba3b49c069c43ee9944 Mon Sep 17 00:00:00 2001 From: Yizhi Liu Date: Sat, 14 Jul 2018 21:33:22 -0700 Subject: [PATCH 7/9] modify doc and private var --- .../src/main/scala/org/apache/mxnet/NDArrayCollector.scala | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala index bbbca23d4174..f61b855efe32 100644 --- a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala +++ b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala @@ -41,7 +41,7 @@ import scala.collection.mutable * * In the case above, the intermediate NDArrays * (created by NDArray.relu and +) will be disposed automatically.
- * User can also decide to use dispose the collected NDArrays later:
+ * User can also decide to dispose the collected NDArrays later:
*
  *  val collector = NDArrayCollector.manual()
  *  val res = collector.withScope {
@@ -92,7 +92,7 @@ object NDArrayCollector {
 
 class NDArrayCollector private(private val autoDispose: Boolean = true,
                                private val doCollect: Boolean = true) {
-  private val arrays: mutable.Map[Long, NDArray] = mutable.HashMap.empty[Long, NDArray]
+  private val arrays = mutable.HashMap.empty[Long, NDArray]
 
   private def add(nd: NDArray*): Unit = {
     if (doCollect) nd.foreach(arr => arrays.put(arr.handle, arr))

From 092268ed46d35e1e9d851f2e5a633e3205e6cd49 Mon Sep 17 00:00:00 2001
From: Yizhi Liu 
Date: Sat, 14 Jul 2018 22:10:22 -0700
Subject: [PATCH 8/9] dispose res when test finishes

---
 .../test/scala/org/apache/mxnet/NDArrayCollectorSuite.scala   | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/scala-package/core/src/test/scala/org/apache/mxnet/NDArrayCollectorSuite.scala b/scala-package/core/src/test/scala/org/apache/mxnet/NDArrayCollectorSuite.scala
index 55fa26b2b0c1..f361ee1e4eac 100644
--- a/scala-package/core/src/test/scala/org/apache/mxnet/NDArrayCollectorSuite.scala
+++ b/scala-package/core/src/test/scala/org/apache/mxnet/NDArrayCollectorSuite.scala
@@ -36,6 +36,8 @@ class NDArrayCollectorSuite extends FunSuite with BeforeAndAfterAll with Matcher
     assert(!res.isDisposed) // smart enough not to dispose the returned NDArray
 
     assert(res.toArray === Array(-1f, 0f, 2f))
+
+    res.dispose()
   }
 
   test("manually dispose") {
@@ -63,5 +65,7 @@ class NDArrayCollectorSuite extends FunSuite with BeforeAndAfterAll with Matcher
 
     collector.clear()
     assert(collector.size === 0)
+
+    res.dispose()
   }
 }

From 60204e18ac2abd20e6830c2c7b9a7b7f0f44bc38 Mon Sep 17 00:00:00 2001
From: Yizhi Liu 
Date: Mon, 16 Jul 2018 21:32:14 -0700
Subject: [PATCH 9/9] add comments, change variable name

---
 .../org/apache/mxnet/NDArrayCollector.scala    | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala
index f61b855efe32..ea21cff9ebc7 100644
--- a/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala
+++ b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala
@@ -17,6 +17,7 @@
 
 package org.apache.mxnet
 
+import org.apache.mxnet.Base.CPtrAddress
 import org.slf4j.LoggerFactory
 
 import scala.annotation.varargs
@@ -92,7 +93,10 @@ object NDArrayCollector {
 
 class NDArrayCollector private(private val autoDispose: Boolean = true,
                                private val doCollect: Boolean = true) {
-  private val arrays = mutable.HashMap.empty[Long, NDArray]
+  // native ptr (handle) of the NDArray -> NDArray
+  // in some rare situation, multiple NDArrays have same native ptr,
+  // the Map here is to prevent from disposing more than once.
+  private val arrays = mutable.HashMap.empty[CPtrAddress, NDArray]
 
   private def add(nd: NDArray*): Unit = {
     if (doCollect) nd.foreach(arr => arrays.put(arr.handle, arr))
@@ -122,20 +126,20 @@ class NDArrayCollector private(private val autoDispose: Boolean = true,
   /**
    * Create a code scope, NDArrays allocated within this scope will be collected.
    * The collected NDArrays will be either 
- * - disposed automatically when the code blcok finishes (when using auto) or
+ * - disposed automatically when the code block finishes (when using auto) or
* - stored for later access (when using manual)
* If the return type of scope is NDArray or NDArrayFuncReturn, * it is smart enough NOT to collect or dispose the returned NDArray.
* However in other cases, it is users' responsibility NOT to leak allocated NDArrays outside. - * @param body code block to be executed within the scope. - * @tparam T return type of the function body. - * @return The result of function body. + * @param codeBlock code block to be executed within the scope. + * @tparam T return type of the function codeBlock. + * @return The result of function codeBlock. */ - def withScope[T](body: => T): T = { + def withScope[T](codeBlock: => T): T = { val old = NDArrayCollector.currCollector.get() NDArrayCollector.currCollector.set(this) try { - val ret = body + val ret = codeBlock ret match { case ndRet: NDArray => arrays.remove(ndRet.handle)