Incubator-mxnet: Java NDArray Disposal

[joshuaswaney]

Description

It appears that neither ResourceScope nor NDArrayCollector work fully as intended in Java.

Following the ResourceScope documentation, it correctly disposes of NDArray objects created within the block but not for the ones returned by Predictor.predictWithNDArray:

NDArray input;
NDArray result;
Predictor predictor;

try (ResourceScope r = new ResourceScope()) {
  input = NDArray.array(...);
  List<NDArray> results = predictor.predictWithNDArray(Arrays.asList(input));
  result = results.get(0);
}
input.nd().isDisposed(); // true
input.nd().isDeAllocated(); // true
result.nd().isDisposed(); // false
result.nd().isDeAllocated(); // false

Following the NDArrayCollector documentation, it does not dispose of any NDArray objects. It looks like NDArrayCollector is deprecated but it's not clear that this is the reason.

NDArray[] outerResult = NDArrayCollector.auto().withScope(new scala.runtime.AbstractFunction0<Object>() {
  @Override
  public Object apply() {
    NDArray input = NDArray.array(...);
    List<NDArray> results = predictor.predictWithNDArray(Arrays.asList(input));
    NDArray result = results.get(0);
    return new NDArray[]{input, result};
  }
});
outerResult[0].nd().isDisposed(); // false
outerResult[0].nd().isDeAllocated(); // false
outerResult[1].nd().isDisposed(); // false
outerResult[1].nd().isDeAllocated(); // false

What have you tried to solve it?

I'm able to work around this by manually disposing NDArray objects using dispose() on each one.

[joshuaswaney]

The wiki says that ResourceScope is "not tested on Java": https://cwiki.apache.org/confluence/display/MXNET/JVM+Memory+Management#JVMMemoryManagement-Java

There's also an open question listed there of: "Will there be a situation where a native pointer is still in use but the Scala Object is not reachable." It appears that this may be the case.

[stu1130]

@joshuaswaney If you are using MXNet java language binding, I would recommend to go with DJL which fully supports MXNet inference and training with CV models. We also have better memory management in NDManager