Repository: deepjavalibrary/djl (java)

Issue:
pytorch doesn't clear out gradients between `backward()` calls in the gradient collector
I've been following along the linear regression code in the Dive into Deep Learning book. When using the pytorch engine the gradients are not cleared between backward propagation runs.

For example (this code is scala but that makes no difference):

```scala
val x = mn.arange(4.0f)
x.setRequiresGradient(true)
println(s"x=$x")

for(i <- 0 until 4) {
  Using.resource(Engine.getInstance().newGradientCollector()){ gc =>
      val y = x.dot(x) * 2
      gc.backward(y)
  }
  println(s"gradient=${x.getGradient()}")
}
```

The following code shows:
```
x=ND: (4) cpu() float32 hasGradient
[0., 1., 2., 3.]

gradient=ND: (4) cpu() float32
[ 0.,  4.,  8., 12.]

gradient=ND: (4) cpu() float32
[ 0.,  8., 16., 24.]

gradient=ND: (4) cpu() float32
[ 0., 12., 24., 36.]

gradient=ND: (4) cpu() float32
[ 0., 16., 32., 48.]
```

ie the gradients keep getting summed between runs when using the pytorch engine but the correct output (below) when swapping out the pytorch dependency for the mxnet one (identical code)
```
x=ND: (4) cpu() float32 hasGradient
[0., 1., 2., 3.]

gradient=ND: (4) cpu() float32
[ 0.,  4.,  8., 12.]

gradient=ND: (4) cpu() float32
[ 0.,  4.,  8., 12.]

gradient=ND: (4) cpu() float32
[ 0.,  4.,  8., 12.]

gradient=ND: (4) cpu() float32
[ 0.,  4.,  8., 12.]
```

This is under MacOS both x86 and ARM (CPU).


Pull request description:
Ensure GradientCollector can clear gradients
It also adds a new function GradientCollector.zeroGradients().

As part of this, it adds a new feature to get the arrays managed by an NDManager or an NDResource. This is so all arrays can be found and have their gradients cleared when the PtGradientCollector is started.

It also means that the System NDManagers are modified to now begin tracking resources. As most things created under them are BaseNDManagers, it shouldn't be a big issue. Then, operations that don't make sense under the SystemNDManagers such as tempAttach or close have been changed to throw an exception rather than work silently.

fixes #2024

Patch:
diff --git a/api/src/main/java/ai/djl/ndarray/BaseNDManager.java b/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
index a26f0c09c..863f2160e 100644
--- a/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/BaseNDManager.java
@@ -32,9 +32,12 @@ import java.nio.LongBuffer;
 import java.nio.ShortBuffer;
 import java.nio.charset.Charset;
 import java.nio.file.Path;
+import java.util.List;
 import java.util.UUID;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 /** {@code BaseNDManager} is the default implementation of {@link NDManager}. */
 public abstract class BaseNDManager implements NDManager {
@@ -298,6 +301,30 @@ public abstract class BaseNDManager implements NDManager {
         return device;
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public List<NDArray> getManagedArrays() {
+        return Stream.concat(
+                        // Main resources
+                        resources.values().stream()
+                                .flatMap(
+                                        r -> {
+                                            if (r instanceof NDResource) {
+                                                return ((NDResource) r)
+                                                        .getResourceNDArrays().stream();
+                                            } else if (r instanceof NDManager) {
+                                                return ((NDManager) r).getManagedArrays().stream();
+                                            } else {
+                                                return Stream.empty();
+                                            }
+                                        }),
+
+                        // Temp resouces
+                        tempResources.values().stream()
+                                .flatMap(tr -> tr.resource.getResourceNDArrays().stream()))
+                .collect(Collectors.toList());
+    }
+
     /** {@inheritDoc} */
     @Override
     public String toString() {
@@ -315,9 +342,6 @@ public abstract class BaseNDManager implements NDManager {
     /** {@inheritDoc} */
     @Override
     public synchronized void attachInternal(String resourceId, AutoCloseable resource) {
-        if (this instanceof SystemNDManager) {
-            return;
-        }
         if (capped.get()) {
             throw new IllegalStateException("NDManager is capped for addition of resources.");
         }
@@ -327,9 +351,6 @@ public abstract class BaseNDManager implements NDManager {
     /** {@inheritDoc} */
     @Override
     public synchronized void attachUncappedInternal(String resourceId, AutoCloseable resource) {
-        if (this instanceof SystemNDManager) {
-            return;
-        }
         if (closed.get()) {
             throw new IllegalStateException("NDManager has been closed already.");
         }
@@ -356,7 +377,8 @@ public abstract class BaseNDManager implements NDManager {
     public void tempAttachInternal(
             NDManager originalManager, String resourceId, NDResource resource) {
         if (this instanceof SystemNDManager) {
-            return;
+            throw new IllegalStateException(
+                    "System manager cannot be temp attached because it can't be closed..");
         }
         if (closed.get()) {
             throw new IllegalStateException("NDManager has been closed already.");
@@ -367,9 +389,6 @@ public abstract class BaseNDManager implements NDManager {
     /** {@inheritDoc} */
     @Override
     public synchronized void detachInternal(String resourceId) {
-        if (this instanceof SystemNDManager) {
-            return;
-        }
         if (closed.get()) {
             // This may happen in the middle of BaseNDManager.close()
             return;
@@ -400,7 +419,9 @@ public abstract class BaseNDManager implements NDManager {
     @Override
     public void close() {
         if (this instanceof SystemNDManager) {
-            return;
+            throw new IllegalStateException(
+                    "The SystemNDManager can not be closed. It is global and lives for the duration"
+                            + " of the process");
         }
         if (!closed.getAndSet(true)) {
             for (AutoCloseable closeable : resources.values()) {
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index f8d1427c8..c7509b268 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -30,6 +30,8 @@ import java.nio.LongBuffer;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import java.util.function.Function;
 import java.util.stream.IntStream;
 import java.util.stream.LongStream;
@@ -4567,6 +4569,12 @@ public interface NDArray extends NDResource, BytesSupplier {
      */
     NDArray erfinv();
 
+    /** {@inheritDoc} */
+    @Override
+    default List<NDArray> getResourceNDArrays() {
+        return Collections.singletonList(this);
+    }
+
     /**
      * Returns an internal representative of Native {@code NDArray}.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/NDList.java b/api/src/main/java/ai/djl/ndarray/NDList.java
index c88c93337..36c73bb42 100644
--- a/api/src/main/java/ai/djl/ndarray/NDList.java
+++ b/api/src/main/java/ai/djl/ndarray/NDList.java
@@ -27,6 +27,7 @@ import java.nio.ByteBuffer;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.List;
 import java.util.zip.ZipEntry;
 import java.util.zip.ZipInputStream;
 import java.util.zip.ZipOutputStream;
@@ -269,6 +270,12 @@ public class NDList extends ArrayList<NDArray> implements NDResource, BytesSuppl
         return head().getManager();
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public List<NDArray> getResourceNDArrays() {
+        return this;
+    }
+
     /** {@inheritDoc} */
     @Override
     public void attach(NDManager manager) {
diff --git a/api/src/main/java/ai/djl/ndarray/NDManager.java b/api/src/main/java/ai/djl/ndarray/NDManager.java
index 0426a16c2..72d4c08e1 100644
--- a/api/src/main/java/ai/djl/ndarray/NDManager.java
+++ b/api/src/main/java/ai/djl/ndarray/NDManager.java
@@ -34,6 +34,7 @@ import java.nio.LongBuffer;
 import java.nio.charset.Charset;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
+import java.util.List;
 
 /**
  * NDArray managers are used to create <I>NDArrays</I> (n-dimensional array on native engine).
@@ -1497,6 +1498,13 @@ public interface NDManager extends AutoCloseable {
      */
     Device getDevice();
 
+    /**
+     * Returns all {@link NDArray}s managed by this manager (including recursively).
+     *
+     * @return all {@link NDArray}s managed by this manager (including recursively)
+     */
+    List<NDArray> getManagedArrays();
+
     /**
      * Attaches a resource to this {@code NDManager}.
      *
diff --git a/api/src/main/java/ai/djl/ndarray/NDResource.java b/api/src/main/java/ai/djl/ndarray/NDResource.java
index 72d709c8d..02d09bb77 100644
--- a/api/src/main/java/ai/djl/ndarray/NDResource.java
+++ b/api/src/main/java/ai/djl/ndarray/NDResource.java
@@ -12,6 +12,8 @@
  */
 package ai.djl.ndarray;
 
+import java.util.List;
+
 /** An object which is managed by an {@link NDManager} and tracks the manager it is attached to. */
 public interface NDResource extends AutoCloseable {
 
@@ -22,6 +24,13 @@ public interface NDResource extends AutoCloseable {
      */
     NDManager getManager();
 
+    /**
+     * Returns the {@link NDArray} or {@link NDArray}s contained within this resource.
+     *
+     * @return the {@link NDArray} or {@link NDArray}s contained within this resource
+     */
+    List<NDArray> getResourceNDArrays();
+
     /**
      * Attaches this {@link NDResource} to the specified {@link NDManager}.
      *
diff --git a/api/src/main/java/ai/djl/training/GradientCollector.java b/api/src/main/java/ai/djl/training/GradientCollector.java
index e1a3cd99c..edadf2158 100644
--- a/api/src/main/java/ai/djl/training/GradientCollector.java
+++ b/api/src/main/java/ai/djl/training/GradientCollector.java
@@ -31,6 +31,9 @@ public interface GradientCollector extends AutoCloseable {
      */
     void backward(NDArray target);
 
+    /** Sets all the gradients within the engine to zero. */
+    void zeroGradients();
+
     /** {@inheritDoc} */
     @Override
     void close();
diff --git a/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java b/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java
index a1cfd9b87..1a0ac0d4e 100644
--- a/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java
+++ b/api/src/main/java/ai/djl/util/passthrough/PassthroughNDManager.java
@@ -26,6 +26,8 @@ import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.charset.Charset;
 import java.nio.file.Path;
+import java.util.Collections;
+import java.util.List;
 
 /** An {@link NDManager} that does nothing, for use in extensions and hybrid engines. */
 public final class PassthroughNDManager implements NDManager {
@@ -241,6 +243,12 @@ public final class PassthroughNDManager implements NDManager {
         return Device.cpu();
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public List<NDArray> getManagedArrays() {
+        return Collections.emptyList();
+    }
+
     /** {@inheritDoc} */
     @Override
     public void attachInternal(String resourceId, AutoCloseable resource) {
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxGradientCollector.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxGradientCollector.java
index 5e1ad8369..a8a183de3 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxGradientCollector.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxGradientCollector.java
@@ -19,7 +19,7 @@ import ai.djl.ndarray.NDManager;
 import ai.djl.training.GradientCollector;
 
 /** {@code MxGradientCollector} is the MXNet implementation of {@link GradientCollector}. */
-public class MxGradientCollector implements GradientCollector {
+public final class MxGradientCollector implements GradientCollector {
 
     /**
      * Constructs an {@code MxGradientCollector} and enables training data collection for
@@ -116,4 +116,15 @@ public class MxGradientCollector implements GradientCollector {
     private void backward(NDArray array, boolean retainGraph) {
         JnaUtils.autogradBackward(new NDList(array), retainGraph ? 1 : 0);
     }
+
+    /** {@inheritDoc} */
+    @Override
+    public void zeroGradients() {
+        NDManager systemManager = MxNDManager.getSystemManager();
+        for (NDArray array : systemManager.getManagedArrays()) {
+            if (array.hasGradient()) {
+                array.getGradient().subi(array.getGradient());
+            }
+        }
+    }
 }
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java
index fad7755fa..e666237f0 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java
@@ -14,11 +14,12 @@
 package ai.djl.pytorch.engine;
 
 import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
 import ai.djl.pytorch.jni.JniUtils;
 import ai.djl.training.GradientCollector;
 
 /** {@code PtGradientCollector} is the PyTorch implementation of {@link GradientCollector}. */
-public class PtGradientCollector implements GradientCollector {
+public final class PtGradientCollector implements GradientCollector {
 
     private boolean gradModel;
 
@@ -26,6 +27,7 @@ public class PtGradientCollector implements GradientCollector {
     public PtGradientCollector() {
         gradModel = JniUtils.isGradMode();
         JniUtils.setGradMode(true);
+        zeroGradients();
     }
 
     /** {@inheritDoc} */
@@ -54,6 +56,17 @@ public class PtGradientCollector implements GradientCollector {
         JniUtils.backward((PtNDArray) target, (PtNDArray) grad, keepGraph, createGraph);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void zeroGradients() {
+        NDManager systemManager = PtNDManager.getSystemManager();
+        for (NDArray array : systemManager.getManagedArrays()) {
+            if (array.hasGradient()) {
+                array.getGradient().subi(array.getGradient());
+            }
+        }
+    }
+
     /** {@inheritDoc} */
     @Override
     public void close() {


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java b/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java
index 340c856f0..1baf9dd37 100644
--- a/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/training/GradientCollectorIntegrationTest.java
@@ -14,6 +14,7 @@ package ai.djl.integration.tests.training;
 
 import ai.djl.Model;
 import ai.djl.basicmodelzoo.basic.Mlp;
+import ai.djl.engine.Engine;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDArrays;
 import ai.djl.ndarray.NDList;
@@ -80,6 +81,46 @@ public class GradientCollectorIntegrationTest {
         }
     }
 
+    @Test
+    public void testZeroGradients() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray a = manager.create(0.0f);
+            a.setRequiresGradient(true);
+
+            try (GradientCollector gc = Engine.getInstance().newGradientCollector()) {
+                NDArray b = a.mul(2);
+
+                // Gradients are initially zero
+                Assert.assertEquals(a.getGradient().getFloat(), 0.0f);
+
+                // Gradients are updated by backwards
+                gc.backward(b);
+                Assert.assertEquals(a.getGradient().getFloat(), 2.0f);
+
+                // Gradients are cleared by zeroGradients
+                gc.zeroGradients();
+                Assert.assertEquals(a.getGradient().getFloat(), 0.0f);
+            }
+        }
+    }
+
+    /** Tests that the gradients do not accumulate when closing the gradient collector. */
+    @Test
+    public void testClearGradients() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray a = manager.create(0.0f);
+            a.setRequiresGradient(true);
+
+            for (int i = 0; i < 3; i++) {
+                try (GradientCollector gc = Engine.getInstance().newGradientCollector()) {
+                    NDArray b = a.mul(2);
+                    gc.backward(b);
+                }
+                Assert.assertEquals(a.getGradient().getFloat(), 2.0f);
+            }
+        }
+    }
+
     @Test
     public void testFreezeParameters() {
         try (Model model = Model.newInstance("model")) {