Repository: deepjavalibrary/djl
Pull request #1943 (2022-08-23 23:33:02)

Adds support for intermediate sequential block results
This makes it so that SequentialBlocks can be used to get the internal results
required for models such as yolo.

codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1943?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1943](https://codecov.io/gh/deepjavalibrary/djl/pull/1943?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (fafb214) into [master](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (bb5073f) will **decrease** coverage by `2.18%`.
> The diff coverage is `67.25%`.

```diff
@@             Coverage Diff              @@
##             master    #1943      +/-   ##
============================================
- Coverage     72.08%   69.90%   -2.19%     
- Complexity     5126     5706     +580     
============================================
  Files           473      544      +71     
  Lines         21970    25178    +3208     
  Branches       2351     2733     +382     
============================================
+ Hits          15838    17600    +1762     
- Misses         4925     6238    +1313     
- Partials       1207     1340     +133     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1943?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1943/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1943/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1943/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1943/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...main/java/ai/djl/modality/cv/output/Rectangle.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1943/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9SZWN0YW5nbGUuamF2YQ==) | `72.41% <0.00%> (ø)` | |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1943/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [...odality/cv/translator/BigGANTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1943/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `33.33% <0.00%> (+8.33%)` | :arrow_up: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1943/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `14.28% <0.00%> (-3.90%)` | :arrow_down: |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1943/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU2VtYW50aWNTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../cv/translator/StyleTransferTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1943/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU3R5bGVUcmFuc2ZlclRyYW5zbGF0b3JGYWN0b3J5LmphdmE=) | `40.00% <ø> (ø)` | |
| ... and [446 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1943/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)


Patch:
diff --git a/api/src/main/java/ai/djl/nn/SequentialBlock.java b/api/src/main/java/ai/djl/nn/SequentialBlock.java
index 7c84fa72d..dae286fce 100644
--- a/api/src/main/java/ai/djl/nn/SequentialBlock.java
+++ b/api/src/main/java/ai/djl/nn/SequentialBlock.java
@@ -22,10 +22,14 @@ import ai.djl.training.ParameterStore;
 import ai.djl.util.PairList;
 
 import java.io.DataInputStream;
+import java.io.DataOutputStream;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.List;
 import java.util.function.Function;
+import java.util.stream.Collectors;
 
 /**
  * {@code SequentialBlock} is a {@link Block} whose children form a chain of blocks with each child
@@ -36,7 +40,8 @@ import java.util.function.Function;
  */
 public class SequentialBlock extends AbstractBlock {
 
-    private static final byte VERSION = 2;
+    private static final byte VERSION = 3;
+    private boolean returnIntermediate;
 
     /**
      * Creates an empty sequential block. Use {@code add} and {@code addAll} to add blocks to be
@@ -148,6 +153,29 @@ public class SequentialBlock extends AbstractBlock {
         }
     }
 
+    /**
+     * Returns whether the block returns all intermediate block results or only the end of the
+     * sequential chain.
+     *
+     * @return whether the block returns all intermediate block results or only the end of the
+     *     sequential chain
+     */
+    public boolean isReturnIntermediate() {
+        return returnIntermediate;
+    }
+
+    /**
+     * Sets whether the block returns all intermediate sequence results.
+     *
+     * @param returnIntermediate true for intermediates, false for only chain result (default and
+     *     typical behavior is false)
+     * @return this {@link SequentialBlock}
+     */
+    public SequentialBlock setReturnIntermediate(boolean returnIntermediate) {
+        this.returnIntermediate = returnIntermediate;
+        return this;
+    }
+
     /** {@inheritDoc} */
     @Override
     protected NDList forwardInternal(
@@ -155,9 +183,15 @@ public class SequentialBlock extends AbstractBlock {
             NDList inputs,
             boolean training,
             PairList<String, Object> params) {
+        List<NDList> past = new ArrayList<>(children.size());
         NDList current = inputs;
         for (Block block : children.values()) {
             current = block.forward(parameterStore, current, training);
+            past.add(current);
+        }
+        if (returnIntermediate) {
+            return new NDList(
+                    past.stream().flatMap(Collection::stream).collect(Collectors.toList()));
         }
         return current;
     }
@@ -169,9 +203,15 @@ public class SequentialBlock extends AbstractBlock {
             NDList data,
             NDList labels,
             PairList<String, Object> params) {
+        List<NDList> past = new ArrayList<>(children.size());
         NDList current = data;
         for (Block block : children.values()) {
             current = block.forward(parameterStore, current, labels, params);
+            past.add(current);
+        }
+        if (returnIntermediate) {
+            return new NDList(
+                    past.stream().flatMap(Collection::stream).collect(Collectors.toList()));
         }
         return current;
     }
@@ -192,20 +232,35 @@ public class SequentialBlock extends AbstractBlock {
         if (children.isEmpty()) {
             throw new IllegalArgumentException("The sequential block is empty");
         }
+        List<Shape[]> past = new ArrayList<>(children.size());
         Shape[] current = inputs;
         for (Block block : children.values()) {
             current = block.getOutputShapes(current);
+            past.add(current);
+        }
+        if (returnIntermediate) {
+            return past.stream().flatMap(Arrays::stream).toArray(Shape[]::new);
         }
         return current;
     }
 
+    /** {@inheritDoc} */
+    @Override
+    protected void saveMetadata(DataOutputStream os) throws IOException {
+        saveInputShapes(os);
+        os.writeBoolean(returnIntermediate);
+    }
+
     /** {@inheritDoc} */
     @Override
     public void loadMetadata(byte loadVersion, DataInputStream is)
             throws IOException, MalformedModelException {
         if (loadVersion == version) {
             readInputShapes(is);
-        } else if (loadVersion != 1) {
+            returnIntermediate = is.readBoolean();
+        } else if (loadVersion == 2) {
+            readInputShapes(is);
+        } else {
             throw new MalformedModelException("Unsupported encoding version: " + loadVersion);
         }
     }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/Encoding.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/Encoding.java
index dfbee95de..e7c48f504 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/Encoding.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/Encoding.java
@@ -25,7 +25,7 @@ public class Encoding {
     private long[] specialTokenMask;
     private CharSpan[] charTokenSpans;
 
-    Encoding(
+    protected Encoding(
             long[] ids,
             long[] typeIds,
             String[] tokens,
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java
index e0bc80cc2..bbf4e35bd 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/tokenizers/HuggingFaceTokenizer.java
@@ -16,6 +16,7 @@ import ai.djl.huggingface.tokenizers.jni.CharSpan;
 import ai.djl.huggingface.tokenizers.jni.LibUtils;
 import ai.djl.huggingface.tokenizers.jni.TokenizersLibrary;
 import ai.djl.modality.nlp.preprocess.Tokenizer;
+import ai.djl.ndarray.NDManager;
 import ai.djl.translate.ArgumentsUtil;
 import ai.djl.util.NativeResource;
 import ai.djl.util.Utils;
@@ -493,6 +494,7 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
     public static final class Builder {
 
         private Path tokenizerPath;
+        private NDManager manager;
         private Map<String, String> options;
 
         Builder() {
@@ -500,6 +502,17 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
             options.put("addSpecialTokens", "true");
         }
 
+        /**
+         * Sets the optional manager used to manage the lifecycle of the tokenizer.
+         *
+         * @param manager the {@link NDManager}
+         * @return this builder
+         */
+        public Builder optManager(NDManager manager) {
+            this.manager = manager;
+            return this;
+        }
+
         /**
          * Sets the name of the tokenizer.
          *
@@ -619,6 +632,19 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
             }
         }
 
+        /**
+         * Utility to make a tokenizer managed by the builder manager (if one is specified).
+         *
+         * @param tokenizer the tokenizer to manage
+         * @return the updated tokenizer
+         */
+        private HuggingFaceTokenizer managed(HuggingFaceTokenizer tokenizer) {
+            if (manager != null) {
+                manager.attachInternal(tokenizer.getUid(), tokenizer);
+            }
+            return tokenizer;
+        }
+
         /**
          * Builds the translator.
          *
@@ -628,12 +654,12 @@ public final class HuggingFaceTokenizer extends NativeResource<Long> implements
         public HuggingFaceTokenizer build() throws IOException {
             String tokenizerName = options.get("tokenizer");
             if (tokenizerName != null) {
-                return HuggingFaceTokenizer.newInstance(tokenizerName, options);
+                return managed(HuggingFaceTokenizer.newInstance(tokenizerName, options));
             }
             if (tokenizerPath == null) {
                 throw new IllegalArgumentException("Missing tokenizer path.");
             }
-            return HuggingFaceTokenizer.newInstance(tokenizerPath, options);
+            return managed(HuggingFaceTokenizer.newInstance(tokenizerPath, options));
         }
     }
 }
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java
index 0026b6ea9..71daf9488 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/FillMaskTranslatorFactory.java
@@ -54,9 +54,10 @@ public class FillMaskTranslatorFactory implements TranslatorFactory {
         Path modelPath = model.getModelPath();
         try {
             HuggingFaceTokenizer tokenizer =
-                    HuggingFaceTokenizer.builder(arguments).optTokenizerPath(modelPath).build();
-            // close tokenizer when model is closed
-            model.getNDManager().attachInternal(tokenizer.getUid(), tokenizer);
+                    HuggingFaceTokenizer.builder(arguments)
+                            .optTokenizerPath(modelPath)
+                            .optManager(model.getNDManager())
+                            .build();
             FillMaskTranslator translator =
                     FillMaskTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == Classifications.class) {
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslatorFactory.java
index 4a10cb27a..87fbe5a20 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/QuestionAnsweringTranslatorFactory.java
@@ -54,9 +54,10 @@ public class QuestionAnsweringTranslatorFactory implements TranslatorFactory {
         Path modelPath = model.getModelPath();
         try {
             HuggingFaceTokenizer tokenizer =
-                    HuggingFaceTokenizer.builder(arguments).optTokenizerPath(modelPath).build();
-            // close tokenizer when model is closed
-            model.getNDManager().attachInternal(tokenizer.getUid(), tokenizer);
+                    HuggingFaceTokenizer.builder(arguments)
+                            .optTokenizerPath(modelPath)
+                            .optManager(model.getNDManager())
+                            .build();
             QuestionAnsweringTranslator translator =
                     QuestionAnsweringTranslator.builder(tokenizer, arguments).build();
             if (input == QAInput.class && output == String.class) {
diff --git a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java
index fc08812a4..e590bf7f7 100644
--- a/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java
+++ b/extensions/tokenizers/src/main/java/ai/djl/huggingface/translator/TokenClassificationTranslatorFactory.java
@@ -54,9 +54,10 @@ public class TokenClassificationTranslatorFactory implements TranslatorFactory {
         Path modelPath = model.getModelPath();
         try {
             HuggingFaceTokenizer tokenizer =
-                    HuggingFaceTokenizer.builder(arguments).optTokenizerPath(modelPath).build();
-            // close tokenizer when model is closed
-            model.getNDManager().attachInternal(tokenizer.getUid(), tokenizer);
+                    HuggingFaceTokenizer.builder(arguments)
+                            .optTokenizerPath(modelPath)
+                            .optManager(model.getNDManager())
+                            .build();
             TokenClassificationTranslator translator =
                     TokenClassificationTranslator.builder(tokenizer, arguments).build();
             if (input == String.class && output == NamedEntity[].class) {
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/ResNetV1.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/ResNetV1.java
index 471086697..65a5e1c7b 100644
--- a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/ResNetV1.java
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/ResNetV1.java
@@ -170,7 +170,7 @@ public final class ResNetV1 {
      * @param builder the {@link Builder} with the necessary arguments
      * @return a {@link Block} that represents the required ResNet model
      */
-    public static Block resnet(Builder builder) {
+    public static SequentialBlock resnet(Builder builder) {
         int numStages = builder.units.length;
         long height = builder.imageShape.get(1);
         SequentialBlock resNet = new SequentialBlock();
@@ -300,7 +300,7 @@ public final class ResNetV1 {
          *
          * @return the {@link ResNetV1} block
          */
-        public Block build() {
+        public SequentialBlock build() {
             if (imageShape == null) {
                 throw new IllegalArgumentException("Must set imageShape");
             }


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/ResnetTest.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/ResnetTest.java
index 31d46cd4a..c757c8dd3 100644
--- a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/ResnetTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/ResnetTest.java
@@ -44,6 +44,7 @@ import ai.djl.translate.NoopTranslator;
 import ai.djl.translate.TranslateException;
 import ai.djl.util.PairList;
 
+import org.testng.Assert;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
@@ -102,6 +103,40 @@ public class ResnetTest {
         }
     }
 
+    @Test
+    public void testWithIntermediate() throws TranslateException {
+        TrainingConfig config =
+                new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
+                        .optDevices(Engine.getInstance().getDevices(2))
+                        .optInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+
+        Block resNet50 =
+                ResNetV1.builder()
+                        .setImageShape(new Shape(1, 28, 28))
+                        .setNumLayers(50)
+                        .setOutSize(10)
+                        .build()
+                        .setReturnIntermediate(true);
+
+        try (Model model = Model.newInstance("resnet")) {
+            model.setBlock(resNet50);
+            int batchSize = 1;
+            Shape inputShape = new Shape(batchSize, 1, 28, 28);
+            int resultSize = 29;
+            try (Trainer trainer = model.newTrainer(config)) {
+                trainer.initialize(inputShape);
+
+                Assert.assertEquals(
+                        resNet50.getOutputShapes(new Shape[] {inputShape}).length, resultSize);
+            }
+            try (Predictor<NDList, NDList> predictor = model.newPredictor(new NoopTranslator())) {
+                NDManager manager = model.getNDManager();
+                NDList result = predictor.predict(new NDList(manager.ones(inputShape)));
+                Assert.assertEquals(result.size(), resultSize);
+            }
+        }
+    }
+
     @Test
     public void testLoadPredict()
             throws IOException, ModelNotFoundException, TranslateException,