Repository: deepjavalibrary/djl
Pull request #2102 (2022-10-25 15:49:13)

Update docs top level memu
## Description ##

Brief description of what this PR is about

- If this change is a backward incompatible change, why must this change be made?
- Interesting edge cases to note here



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/docs/mkdocs.yml b/docs/mkdocs.yml
index 4c31f4de3..c0f231d28 100644
--- a/docs/mkdocs.yml
+++ b/docs/mkdocs.yml
@@ -89,47 +89,49 @@ nav:
   - DJL Community:
       - 'docs/forums.md'
       - 'leaders.md'
-  - Apache MXNet:
-      - Overview: 'engines/mxnet/README.md'
-      - MXNet Engine: 'engines/mxnet/mxnet-engine/README.md'
-      - MXNet Model Zoo: 'engines/mxnet/mxnet-model-zoo/README.md'
-      - Import Gluon Model: 'docs/mxnet/how_to_convert_your_model_to_symbol.md'
-      - Load a MXNet Model: 'jupyter/load_mxnet_model.ipynb'
-      - Backend Optimizer for MXNet: 'docs/mxnet/mxnet_backend_optimizer.md'
-  - PyTorch:
-      - Overview: 'engines/pytorch/README.md'
-      - PyTorch Engine: 'engines/pytorch/pytorch-engine/README.md'
-      - PyTorch Model Zoo: 'engines/pytorch/pytorch-model-zoo/README.md'
-      - Import PyTorch Model: 'docs/pytorch/how_to_convert_your_model_to_torchscript.md'
-      - Load a PyTorch Model: 'jupyter/load_pytorch_model.ipynb'
-  - TensorFlow:
-      - Overview: 'engines/tensorflow/README.md'
-      - TensorFlow Engine: 'engines/tensorflow/tensorflow-engine/README.md'
-      - TensorFlow Model Zoo: 'engines/tensorflow/tensorflow-model-zoo/README.md'
-      - Import TensorFlow Model: 'docs/tensorflow/how_to_import_tensorflow_models_in_DJL.md'
-      - Load a TensorFlow Model: 'jupyter/tensorflow/pneumonia_detection.ipynb'
-  - Hybrid engines:
-      - Hybrid engine overview: 'docs/hybrid_engine.md'
-      - ONNX Runtime:
-          - Overview: 'engines/onnxruntime/onnxruntime-engine/README.md'
-          - Load a ONNX Model: 'jupyter/onnxruntime/machine_learning_with_ONNXRuntime.ipynb'
-      - PaddlePaddle:
-          - Overview: 'engines/paddlepaddle/README.md'
-          - PaddlePaddle Engine: 'engines/paddlepaddle/paddlepaddle-engine/README.md'
-          - PaddlePaddle Model Zoo: 'engines/paddlepaddle/paddlepaddle-model-zoo/README.md'
-          - Import Paddle model:
-              - English: 'docs/paddlepaddle/how_to_create_paddlepaddle_model.md'
-              - 中文: 'docs/paddlepaddle/how_to_create_paddlepaddle_model_zh.md'
-          - Facemask detection using PaddlePaddle:
-              - English: 'jupyter/paddlepaddle/face_mask_detection_paddlepaddle.ipynb'
-              - 中文: 'jupyter/paddlepaddle/face_mask_detection_paddlepaddle_zh.ipynb'
-          - PaddleOCR example:
-              - English: 'jupyter/paddlepaddle/paddle_ocr_java.ipynb'
-              - 中文: 'jupyter/paddlepaddle/paddle_ocr_java_zh.ipynb'
-      - TensorRT: 'engines/tensorrt/README.md'
-      - TensorFlow Lite: 'engines/tflite/tflite-engine/README.md'
-      - Neo DLR: 'engines/dlr/dlr-engine/README.md'
-      - XGBoost: 'engines/ml/xgboost/README.md'
+  - Supported Engines:
+    - PyTorch:
+        - Overview: 'engines/pytorch/README.md'
+        - PyTorch Engine: 'engines/pytorch/pytorch-engine/README.md'
+        - PyTorch Model Zoo: 'engines/pytorch/pytorch-model-zoo/README.md'
+        - Import PyTorch Model: 'docs/pytorch/how_to_convert_your_model_to_torchscript.md'
+        - Load a PyTorch Model: 'jupyter/load_pytorch_model.ipynb'
+    - TensorFlow:
+        - Overview: 'engines/tensorflow/README.md'
+        - TensorFlow Engine: 'engines/tensorflow/tensorflow-engine/README.md'
+        - TensorFlow Model Zoo: 'engines/tensorflow/tensorflow-model-zoo/README.md'
+        - Import TensorFlow Model: 'docs/tensorflow/how_to_import_tensorflow_models_in_DJL.md'
+        - Load a TensorFlow Model: 'jupyter/tensorflow/pneumonia_detection.ipynb'
+    - Apache MXNet:
+        - Overview: 'engines/mxnet/README.md'
+        - MXNet Engine: 'engines/mxnet/mxnet-engine/README.md'
+        - MXNet Model Zoo: 'engines/mxnet/mxnet-model-zoo/README.md'
+        - Import Gluon Model: 'docs/mxnet/how_to_convert_your_model_to_symbol.md'
+        - Load a MXNet Model: 'jupyter/load_mxnet_model.ipynb'
+        - Backend Optimizer for MXNet: 'docs/mxnet/mxnet_backend_optimizer.md'
+    - Hybrid engines:
+        - Hybrid engine overview: 'docs/hybrid_engine.md'
+        - ONNX Runtime:
+            - Overview: 'engines/onnxruntime/onnxruntime-engine/README.md'
+            - Load a ONNX Model: 'jupyter/onnxruntime/machine_learning_with_ONNXRuntime.ipynb'
+        - PaddlePaddle:
+            - Overview: 'engines/paddlepaddle/README.md'
+            - PaddlePaddle Engine: 'engines/paddlepaddle/paddlepaddle-engine/README.md'
+            - PaddlePaddle Model Zoo: 'engines/paddlepaddle/paddlepaddle-model-zoo/README.md'
+            - Import Paddle model:
+                - English: 'docs/paddlepaddle/how_to_create_paddlepaddle_model.md'
+                - 中文: 'docs/paddlepaddle/how_to_create_paddlepaddle_model_zh.md'
+            - Facemask detection using PaddlePaddle:
+                - English: 'jupyter/paddlepaddle/face_mask_detection_paddlepaddle.ipynb'
+                - 中文: 'jupyter/paddlepaddle/face_mask_detection_paddlepaddle_zh.ipynb'
+            - PaddleOCR example:
+                - English: 'jupyter/paddlepaddle/paddle_ocr_java.ipynb'
+                - 中文: 'jupyter/paddlepaddle/paddle_ocr_java_zh.ipynb'
+        - XGBoost: 'engines/ml/xgboost/README.md'
+        - LightGBM: 'engines/ml/lightgbm/README.md'
+        - TensorRT: 'engines/tensorrt/README.md'
+        - TensorFlow Lite: 'engines/tflite/tflite-engine/README.md'
+        - Neo DLR: 'engines/dlr/dlr-engine/README.md'
   - Extensions:
       - Android: 'android/README.md'
       - AWS S3 support: 'extensions/aws-ai/README.md'
@@ -140,6 +142,18 @@ nav:
       - SentencePiece: 'extensions/sentencepiece/README.md'
       - Tablesaw: 'extensions/tablesaw/README.md'
       - DJL Zero: 'djl-zero/README.md'
+      - DJL Benchmark: 'docs/serving/benchmark/README.md'
+  - DJL Serving:
+      - Why DJL Serving?: 'docs/serving/README.md'
+      - 'docs/serving/serving/docs/starting.md'
+      - 'docs/serving/serving/docs/configuration.md'
+      - 'docs/serving/serving/docs/workflows.md'
+      - 'docs/serving/serving/docs/architecture.md'
+      - HTTP API:
+          - 'docs/serving/serving/docs/inference_api.md'
+          - 'docs/serving/serving/docs/management_api.md'
+      - 'docs/serving/serving/docs/plugin_management.md'
+      - 'docs/serving/wlm/README.md'
   - Demos:
       - Demos: 'docs/demos/README.md'
       - AWS:
@@ -182,16 +196,3 @@ nav:
       - Extensions:
           - Visualizing Training with DJL: 'docs/demos/visualization-vue/README.md'
           - Interactive JShell and Block Runner for DJL: 'docs/demos/web-demo/interactive-console/README.md'
-  - DJL Serving:
-      - Why DJL Serving?: 'docs/serving/README.md'
-      - 'docs/serving/serving/docs/starting.md'
-      - 'docs/serving/serving/docs/configuration.md'
-      - 'docs/serving/serving/docs/workflows.md'
-      - 'docs/serving/serving/docs/architecture.md'
-      - HTTP API:
-        - 'docs/serving/serving/docs/inference_api.md'
-        - 'docs/serving/serving/docs/management_api.md'
-      - 'docs/serving/serving/docs/plugin_management.md'
-      - 'docs/serving/wlm/README.md'
-
-
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")) {