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..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
@@ -554,11 +554,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] = {
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..ea21cff9ebc7
--- /dev/null
+++ b/scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala
@@ -0,0 +1,159 @@
+/*
+ * 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.apache.mxnet.Base.CPtrAddress
+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 code block finishes, or
+ * - simply be collected for future usage.
+ *
+ * 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 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 = {
+ currCollector.get().add(ndArray: _*)
+ }
+}
+
+class NDArrayCollector private(private val autoDispose: Boolean = true,
+ private val doCollect: Boolean = true) {
+ // 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))
+ }
+
+ /**
+ * 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 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 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](codeBlock: => T): T = {
+ val old = NDArrayCollector.currCollector.get()
+ NDArrayCollector.currCollector.set(this)
+ try {
+ val ret = codeBlock
+ ret match {
+ case ndRet: NDArray =>
+ arrays.remove(ndRet.handle)
+ case ndarrays: NDArrayFuncReturn =>
+ ndarrays.arr.foreach(nd => arrays.remove(nd.handle))
+ case _ => // do nothing
+ }
+ ret
+ } finally {
+ if (autoDispose) {
+ foreach(_.dispose())
+ clear()
+ }
+ 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..f361ee1e4eac
--- /dev/null
+++ b/scala-package/core/src/test/scala/org/apache/mxnet/NDArrayCollectorSuite.scala
@@ -0,0 +1,71 @@
+/*
+ * 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))
+
+ res.dispose()
+ }
+
+ 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)
+
+ res.dispose()
+ }
+}
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);