[MXNET-600][Scala] NDArray auto-collector #11751

yzhliu · 2018-07-13T17:49:43Z

For now user has to dispose all intermediate allocated NDArrays themselves. The NDArrayCollector introduced here is to provide a way for NDs to be disposed automatically.

In the test case & changes to NDArrayIter.scala one can find the example of how to use it.

And for Java users:
https://github.com/yzhliu/mxnet-java-example/blob/master/src/main/java/me/yzhi/mxnet/example/infer/RndImageInference.java#L28-L30
https://github.com/yzhliu/mxnet-java-example/blob/master/src/main/java/me/yzhi/mxnet/example/infer/Seq2SeqInference.java#L35-L37

In the cases above, NDArrays allocated within the scope will be disposed automatically after the code block finishes.

lanking520 · 2018-07-13T21:32:11Z

Hi @yzhliu , It's been so great to see this solution. I would recommend you to apply this new feature to the examples in the CI for testing and that include MNIST and GAN to see if there are any more memory leaks thrown.

szha · 2018-07-14T01:24:40Z

@yzhliu nice to see you back on Scala for core development.

CodingCat · 2018-07-15T03:54:17Z

scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala

+}
+
+class NDArrayCollector private(private val autoDispose: Boolean = true,
+ private val doCollect: Boolean = true) {


do we really need this flag, in which case we will set it to false?

https://github.com/apache/incubator-mxnet/pull/11751/files#diff-d502d315f6c6df78673dcde4d27a9577R69
new NDArray always do NDArrayCollector.collect(this), but nothing be really collected unless the user explicitly uses withScope

CodingCat · 2018-07-15T03:55:10Z

scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala

+ * </pre>
+ * In the case above, the intermediate NDArrays
+ * (created by <em>NDArray.relu</em> and <em>+</em>) will be disposed automatically. <br />
+ * User can also decide to use dispose the collected NDArrays later: <br />


dispose is extra?

or use is extra?

CodingCat · 2018-07-15T03:58:37Z

scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala

+ * val collector = NDArrayCollector.manual()
+ * val res = collector.withScope {
+ * (NDArray.relu(a) + a).toArray
+ * }


if the user does not want auto disposal, what's the other benefit withScope can bring to him/her?

if there is no other benefit, we may consider make it simpler as just a NDArray.scope, all NDArray in this scope will be automatically disposed (do not even need manual scope)...and for the other case, the user can just do what they are currently doing

I think the ability to collect and dispose manually is useful, for use cases like,

Users may want to dispose only a subset of the new-allocated NDArrays, e.g., the withScope returns a complicated data structure which contains NDArrays, these NDArrays normally cannot be disposed automatically (and cannot easily be detect by withScope.

Debug/performance analysis.

CodingCat · 2018-07-15T04:03:33Z

scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala

+
+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 variable may not need to have an explicit type

yzhliu · 2018-07-15T17:29:20Z

@lanking520 Here you can find how we can fix the memory leak for MNIST training (using SGD): yzhliu#5

But I prefer to make the PR separate. You can try similar thing for GAN, and with other optimizer.

lanking520 · 2018-07-16T04:10:49Z

scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala

+
+class NDArrayCollector private(private val autoDispose: Boolean = true,
+ private val doCollect: Boolean = true) {
+ private val arrays = mutable.HashMap.empty[Long, NDArray]


Please put some comments in here: The first long represent as the C++ ptr of the NDArray, I think type CPtrAddress = Long will be more clearer in here, it is defined in the base.scala

This HashMap would only keep the ptr as the key, there could be multiple NDArray reference point to the same ptr. Will there be any risks if we only keep one NDArray here (such as we lose reference if A NDArray have different parent of B NDArray but we only kill one of them)? In that case, ArrayBuffer[NDArray] will be better?

this is exactly what I intended to avoid. Two NDArrays with same ptr must be disposed only once.

lanking520 · 2018-07-16T04:11:56Z

scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala

+ /**
+ * Clear the collector.
+ */
+ def clear(): Unit = {


I don't see there is any use cases from outside world, shall we keep it private?

When using manual(), user may want to re-use one collector:

val c = NDArrayCollector.manual() c.withScope { ... } ... c.clear() c.withScope { ... }

lanking520 · 2018-07-16T04:13:38Z

scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala

+ /**
+ * Create a code scope, NDArrays allocated within this scope will be collected.
+ * The collected NDArrays will be either <br />
+ * - disposed automatically when the code blcok finishes (when using <em>auto</em>) or <br />


blcok -> block

lanking520 · 2018-07-16T04:20:16Z

scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala

+ * @return The result of function <em>body</em>.
+ */
+ def withScope[T](body: => T): T = {
+ val old = NDArrayCollector.currCollector.get()


Am I correct?

In every local thread, there could be only 1 currCollector exist.

If there is a cascade NDArrayCollector, we still have one currCollector, but we kill the inner loop result and keep outer one there.

As we talked offline, yes.

lanking520 · 2018-07-16T04:25:10Z

scala-package/core/src/main/scala/org/apache/mxnet/NDArrayCollector.scala

+ * @tparam T return type of the function <em>body</em>.
+ * @return The result of function <em>body</em>.
+ */
+ def withScope[T](body: => T): T = {


Do we always need to define a return NDArray if we want to keep something created from this inner scope? Shall we change the body to some better name? Such as returnValue?

nswamy

I really like this solution, its an elegant implementation of enter/exit blocks like in python. I have concerns of using ThreadLocal which is a red flag to me(IMO) since it takes impacts thread-safety in subtle ways which will be hard to debug, for example in this case if a NDArray in WithScope() in thread-A is passed to another thread-B (think of a producer-consumer problem), if thread-A exits before the object is used in thread-B it could cause problems.

yzhliu · 2018-07-16T20:25:20Z

@nswamy Actually it is not a problem with threadlocal. No matter threadlocal is used or not, the NDArrays within the scope cannot be passed to another thread. It is stated in the Javadoc of NDArrayCollector.

ThreadLocal here is for,
Thread A:

collectorA.withScope {
   val a = new NDArray;
}

Thread B:

collectorB.withScope {
  val b = new NDArray;
}

ThreadLocal can guarantee b stays in NDCollectorB and a stays in NDCollectorA. Otherwise, the scope can be messed up, e.g.,

currCollector = NDCollectorA;  // thread A starts
currCollector = NDCollectorB;  // thread B starts, (but) before a = new NDArray runs
val a = new NDArray; // a will be in NDCollectorB
val b = new NDArray; // b will be also in NDCollectorB
// scope B finishes
a.dispose(); b.dispose(). // scope B dispose all its collected NDs
currCollector = NDCollectorA. // scope B reset the current collector
// but scope A is still running, where NDArray a can be accessed.

lanking520

LGTM

nswamy · 2018-07-19T00:17:35Z

@yzhliu Thanks for elaborating with a detailed example and sorry for the delay in response.
While I understand WithScope is what manages the deallocation of the NDArrays and using it ThreadLocal is what makes managing it possible, nevertheless my concern about thread-safety still remains.
The use-case of producer/consumer is not very uncommon and

val q = new Queue[NDArray] //assume Queue  is thread-safe

val t1 = new Thread() {
@override def run() ={
withScope {
 val a = NDArray.load()
 q.enqueue(a)
}
}
val t2 = new Thread() {
@override def run() = {
b = q.dequeue()
b.reshape()
}
}

This use-case or similar is not hard imagine when you are creating data pipeline and I think we should cautious in introducing new patterns to users that aren't thread-safe.

We discussed offline about using reference-counts, any luck with that?

yzhliu · 2018-07-19T02:00:19Z

@nswamy unfortunately I don't think it is possible to implement within Java user code (unless we leverage finalize) - Users cannot access java heap (in our case, reference number of one object) - the only way is to dump from outside.

Moreover, not necessarily to be multi-threading, even single-thread can have problem:

var a: NDArray
collector.withScope {
    a = new NDArray
}  // a is disposed here
a.toArray() // fail

As I stated, it is users' responsibility not to leak new NDArray outside of scope, otherwise they need to use manual(), and dispose it later.

Every tool has its own limitation, e.g., users are not supposed to use HashMap in multi-threads - otherwise they need to do synchronize themselves.

As you can see, this tool make disposing super easy: https://github.com/yzhliu/mxnet/pull/5/files#diff-19b434167a8c2d81db271f8c47975ff6R43

btw, In your case, users need to do:

val q = new Queue[NDArray] //assume Queue  is thread-safe

val c = NDArrayCollector.manual()

val t1 = new Thread() {
@override def run() ={
c.withScope {
 val a = NDArray.load()
 q.enqueue(a)
}
c.foreach(nd => dispose if nd != a)
}

val t2 = new Thread() {
@override def run() = {
b = q.dequeue()
b.reshape()
}
}

t2.join()
c.foreach(_.dispose())

nswamy · 2018-07-19T02:20:57Z

@yzhliu I understand what you are saying and I think it is given that it doesn't work outside of the scope.
agree every tool has its limitation, I already said that I like this :)

Do you think changing WithScope to WithThreadScope and make it explicit would be useful for users to communicate that it is not thread-safe?

Also, I did not intend my code to be a running version, it was more for illustration of my point.. i'll call it out next time :)

yzhliu · 2018-07-19T02:31:43Z

@nswamy Oh no no, I was not saying your code is wrong or cannot run, I was demonstrating how to change the code a bit to safely use the collector in such situation - my code was not a runnable version either :)

As I said, even single thread can have problem, don't think it is proper to have 'Thread' in the name.

nswamy · 2018-07-19T02:44:04Z

you are right, changing to ThreadScope might again confuse people if the case you showed.

Thanks for the clarifying about the example code, appreciate your patience :)

nswamy

Noting that Thread-safety is a limitation of this solution. We can see what users think of this approach in solving the memory management issue.

yzhliu · 2018-07-19T02:56:12Z

Thanks all people helped to review here.

* [Scala] NDArrayCollector for automatically disposing NDArrays * modify doc for NDArrayCollector * modify the function doc of NDArrayCollector.withScope * remove trivial changes * put dispose in finally * fix jni NDArray signature * modify doc and private var * dispose res when test finishes * add comments, change variable name

yzhliu added 3 commits July 12, 2018 17:24

[Scala] NDArrayCollector for automatically disposing NDArrays

9b15341

modify doc for NDArrayCollector

6ab02a3

modify the function doc of NDArrayCollector.withScope

8addba8

yzhliu requested a review from nswamy as a code owner July 13, 2018 17:49

yzhliu changed the title ~~Scala nd collector~~ [Scala] NDArray auto-collector Jul 13, 2018

yzhliu changed the title ~~[Scala] NDArray auto-collector~~ [MXNET-600][Scala] NDArray auto-collector Jul 13, 2018

remove trivial changes

731066d

yzhliu requested a review from CodingCat July 13, 2018 18:22

put dispose in finally

ddf5daa

yzhliu added 2 commits July 13, 2018 16:09

fix jni NDArray signature

72cf085

Merge remote-tracking branch 'upstream/master' into scala-nd-collector

f167c62

CodingCat reviewed Jul 15, 2018

View reviewed changes

yzhliu added 2 commits July 14, 2018 21:33

modify doc and private var

1417472

dispose res when test finishes

092268e

lanking520 suggested changes Jul 16, 2018

View reviewed changes

nswamy reviewed Jul 16, 2018

View reviewed changes

add comments, change variable name

60204e1

CodingCat approved these changes Jul 18, 2018

View reviewed changes

lanking520 approved these changes Jul 18, 2018

View reviewed changes

nswamy approved these changes Jul 19, 2018

View reviewed changes

yzhliu merged commit 1031fe1 into apache:master Jul 19, 2018

lanking520 mentioned this pull request Jul 31, 2018

FeedForward.scala and NDArrayIter.scala leak memory by not disposing of NDArrays #10436

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[MXNET-600][Scala] NDArray auto-collector #11751

[MXNET-600][Scala] NDArray auto-collector #11751

yzhliu commented Jul 13, 2018 •

edited

Loading

lanking520 commented Jul 13, 2018

szha commented Jul 14, 2018

CodingCat Jul 15, 2018

yzhliu Jul 15, 2018

CodingCat Jul 15, 2018

CodingCat Jul 15, 2018

CodingCat Jul 15, 2018

CodingCat Jul 15, 2018

yzhliu Jul 15, 2018

CodingCat Jul 15, 2018

yzhliu commented Jul 15, 2018

lanking520 Jul 16, 2018

yzhliu Jul 17, 2018

lanking520 Jul 16, 2018

yzhliu Jul 17, 2018

lanking520 Jul 16, 2018

lanking520 Jul 16, 2018

yzhliu Jul 17, 2018

lanking520 Jul 16, 2018

nswamy left a comment •

edited

Loading

yzhliu commented Jul 16, 2018 •

edited

Loading

lanking520 left a comment

nswamy commented Jul 19, 2018

yzhliu commented Jul 19, 2018 •

edited

Loading

nswamy commented Jul 19, 2018

yzhliu commented Jul 19, 2018

nswamy commented Jul 19, 2018

nswamy left a comment

yzhliu commented Jul 19, 2018

[MXNET-600][Scala] NDArray auto-collector #11751

[MXNET-600][Scala] NDArray auto-collector #11751

Conversation

yzhliu commented Jul 13, 2018 • edited Loading

lanking520 commented Jul 13, 2018

szha commented Jul 14, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

yzhliu commented Jul 15, 2018

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

nswamy left a comment • edited Loading

Choose a reason for hiding this comment

yzhliu commented Jul 16, 2018 • edited Loading

lanking520 left a comment

Choose a reason for hiding this comment

nswamy commented Jul 19, 2018

yzhliu commented Jul 19, 2018 • edited Loading

nswamy commented Jul 19, 2018

yzhliu commented Jul 19, 2018

nswamy commented Jul 19, 2018

nswamy left a comment

Choose a reason for hiding this comment

yzhliu commented Jul 19, 2018

yzhliu commented Jul 13, 2018 •

edited

Loading

nswamy left a comment •

edited

Loading

yzhliu commented Jul 16, 2018 •

edited

Loading

yzhliu commented Jul 19, 2018 •

edited

Loading