Fix memory leak in zeroGradients()

api/src/main/java/ai/djl/modality/cv/translator/YoloV5Translator.java

@@ -160,7 +160,7 @@ protected double overlap(double x1, double w1, double x2, double w2) {
         return right - left;
     }
 
-    private DetectedObjects processFromBoxOutput(NDList list) {
+    protected DetectedObjects processFromBoxOutput(NDList list) {
         float[] flattened = list.get(0).toFloatArray();
         ArrayList<IntermediateResult> intermediateResults = new ArrayList<>();
         int sizeClasses = classes.size();
@@ -280,7 +280,7 @@ public YoloV5Translator build() {
         }
     }
 
-    private static final class IntermediateResult {
+    protected static final class IntermediateResult {
 
         /**
          * A sortable score for how good the recognition is relative to others. Higher should be

api/src/main/java/ai/djl/modality/cv/translator/YoloV8Translator.java

@@ -0,0 +1,103 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file 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 ai.djl.modality.cv.translator;
+
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.output.Rectangle;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.types.DataType;
+
+import java.util.ArrayList;
+import java.util.Map;
+
+/**
+ * A translator for YoloV8 models. This was tested with ONNX exported Yolo models. For details check
+ * here: https://github.com/ultralytics/ultralytics
+ */
+public class YoloV8Translator extends YoloV5Translator {
+
+    /**
+     * Constructs an ImageTranslator with the provided builder.
+     *
+     * @param builder the data to build with
+     */
+    protected YoloV8Translator(Builder builder) {
+        super(builder);
+    }
+
+    /**
+     * Creates a builder to build a {@code YoloV8Translator} with specified arguments.
+     *
+     * @param arguments arguments to specify builder options
+     * @return a new builder
+     */
+    public static YoloV8Translator.Builder builder(Map<String, ?> arguments) {
+        YoloV8Translator.Builder builder = new YoloV8Translator.Builder();
+        builder.configPreProcess(arguments);
+        builder.configPostProcess(arguments);
+
+        return builder;
+    }
+
+    @Override
+    protected DetectedObjects processFromBoxOutput(NDList list) {
+        NDArray features4OneImg = list.get(0);
+        int sizeClasses = classes.size();
+        long sizeBoxes = features4OneImg.size(1);
+        ArrayList<IntermediateResult> intermediateResults = new ArrayList<>();
+
+        for (long b = 0; b < sizeBoxes; b++) {
+            float maxClass = 0;
+            int maxIndex = 0;
+            for (int c = 4; c < sizeClasses; c++) {
+                float classProb = features4OneImg.getFloat(c, b);
+                if (classProb > maxClass) {
+                    maxClass = classProb;
+                    maxIndex = c;
+                }
+            }
+
+            if (maxClass > threshold) {
+                float xPos = features4OneImg.getFloat(0, b); // center x
+                float yPos = features4OneImg.getFloat(1, b); // center y
+                float w = features4OneImg.getFloat(2, b);
+                float h = features4OneImg.getFloat(3, b);
+                Rectangle rect =
+                        new Rectangle(Math.max(0, xPos - w / 2), Math.max(0, yPos - h / 2), w, h);
+                intermediateResults.add(
+                        new IntermediateResult(classes.get(maxIndex), maxClass, maxIndex, rect));
+            }
+        }
+
+        return nms(intermediateResults);
+    }
+
+    /** The builder for {@link YoloV8Translator}. */
+    public static class Builder extends YoloV5Translator.Builder {
+        /**
+         * Builds the translator.
+         *
+         * @return the new translator
+         */
+        @Override
+        public YoloV8Translator build() {
+            if (pipeline == null) {
+                addTransform(
+                        array -> array.transpose(2, 0, 1).toType(DataType.FLOAT32, false).div(255));
+            }
+            validate();
+            return new YoloV8Translator(this);
+        }
+    }
+}

api/src/main/java/ai/djl/modality/cv/translator/YoloV8TranslatorFactory.java

@@ -0,0 +1,35 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file 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 ai.djl.modality.cv.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.translate.Translator;
+
+import java.io.Serializable;
+import java.util.Map;
+
+/** A translatorFactory that creates a {@link YoloV8Translator} instance. */
+public class YoloV8TranslatorFactory extends ObjectDetectionTranslatorFactory
+        implements Serializable {
+
+    private static final long serialVersionUID = 1L;
+
+    /** {@inheritDoc} */
+    @Override
+    protected Translator<Image, DetectedObjects> buildBaseTranslator(
+            Model model, Map<String, ?> arguments) {
+        return YoloV8Translator.builder(arguments).build();
+    }
+}

api/src/main/java/ai/djl/nn/Block.java

@@ -313,6 +313,21 @@ default void freezeParameters(boolean freeze) {
         }
     }
 
+    /**
+     * Freezes or unfreezes all parameters inside the block that pass the predicate.
+     *
+     * @param freeze true to mark as frozen rather than unfrozen
+     * @param pred true tests if the parameter should be updated
+     * @see Parameter#freeze(boolean)
+     */
+    default void freezeParameters(boolean freeze, Predicate<Parameter> pred) {
+        for (Parameter parameter : getParameters().values()) {
+            if (pred.test(parameter)) {
+                parameter.freeze(freeze);
+            }
+        }
+    }
+
     /**
      * Validates that actual layout matches the expected layout.
      *

api/src/test/java/ai/djl/modality/cv/translator/YoloV8TranslatorFactoryTest.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2023 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file 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 ai.djl.modality.cv.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.translate.BasicTranslator;
+import ai.djl.translate.Translator;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.io.InputStream;
+import java.net.URL;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.Map;
+
+public class YoloV8TranslatorFactoryTest {
+
+    private YoloV8TranslatorFactory factory;
+
+    @BeforeClass
+    public void setUp() {
+        factory = new YoloV8TranslatorFactory();
+    }
+
+    @Test
+    public void testGetSupportedTypes() {
+        Assert.assertEquals(factory.getSupportedTypes().size(), 5);
+    }
+
+    @Test
+    public void testNewInstance() {
+        Map<String, String> arguments = new HashMap<>();
+        try (Model model = Model.newInstance("test")) {
+            Translator<Image, DetectedObjects> translator1 =
+                    factory.newInstance(Image.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator1 instanceof YoloV8Translator);
+
+            Translator<Path, DetectedObjects> translator2 =
+                    factory.newInstance(Path.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator2 instanceof BasicTranslator);
+
+            Translator<URL, DetectedObjects> translator3 =
+                    factory.newInstance(URL.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator3 instanceof BasicTranslator);
+
+            Translator<InputStream, DetectedObjects> translator4 =
+                    factory.newInstance(InputStream.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator4 instanceof BasicTranslator);
+
+            Translator<Input, Output> translator5 =
+                    factory.newInstance(Input.class, Output.class, model, arguments);
+            Assert.assertTrue(translator5 instanceof ImageServingTranslator);
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(Image.class, Output.class, model, arguments));
+        }
+    }
+}

engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxGradientCollector.java

@@ -123,7 +123,10 @@ public void zeroGradients() {
         NDManager systemManager = MxNDManager.getSystemManager();
         for (NDArray array : systemManager.getManagedArrays()) {
             if (array.hasGradient()) {
-                array.getGradient().subi(array.getGradient());
+                // To prevent memory leak we must close gradient after use.
+                try (NDArray gradient = array.getGradient()) {
+                    gradient.subi(gradient);
+                }
             }
         }
     }

engines/onnxruntime/onnxruntime-engine/src/main/java/ai/djl/onnxruntime/engine/OrtEngine.java

@@ -97,7 +97,7 @@ public int getRank() {
     /** {@inheritDoc} */
     @Override
     public String getVersion() {
-        return "1.15.1";
+        return "1.16.0";
     }
 
     /** {@inheritDoc} */

engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtGradientCollector.java

@@ -76,7 +76,10 @@ public void zeroGradients() {
         NDManager systemManager = PtNDManager.getSystemManager();
         for (NDArray array : systemManager.getManagedArrays()) {
             if (array.hasGradient()) {
-                array.getGradient().subi(array.getGradient());
+                // To prevent memory leak we must close gradient after use.
+                try (NDArray gradient = array.getGradient()) {
+                    gradient.subi(gradient);
+                }
             }
         }
     }

engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtModel.java

@@ -18,6 +18,7 @@
 import ai.djl.Model;
 import ai.djl.ndarray.types.DataType;
 import ai.djl.nn.Parameter;
+import ai.djl.nn.Parameter.Type;
 import ai.djl.pytorch.jni.JniUtils;
 import ai.djl.training.Trainer;
 import ai.djl.training.TrainingConfig;
@@ -189,7 +190,9 @@ public Trainer newTrainer(TrainingConfig trainingConfig) {
         }
         if (wasLoaded) {
             // Unfreeze parameters if training directly
-            block.freezeParameters(false);
+            block.freezeParameters(
+                    false,
+                    p -> p.getType() != Type.RUNNING_MEAN && p.getType() != Type.RUNNING_VAR);
         }
         for (Pair<Initializer, Predicate<Parameter>> pair : initializer) {
             if (pair.getKey() != null && pair.getValue() != null) {