Repository: deepjavalibrary/djl
Pull request #2981 (2024-02-06 19:40:52)

Supports PyTorch stream imperative model load
This adds support for PyTorch InputStream model load when using imperative models. Before, it would assume symbolic without actually checking. In those cases, it would throw an engine exception. Now, that option is working and tested.

codecov-commenter: ## [Codecov](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Attention: `1330 lines` in your changes are missing coverage. Please review.
> Comparison is base [(`bb5073f`)](https://app.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) 72.08% compared to head [(`681af48`)](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) 72.10%.
> Report is 986 commits behind head on master.

| [Files](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Patch % | Lines |
|---|---|---|
| [...va/ai/djl/modality/nlp/generate/TextGenerator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9UZXh0R2VuZXJhdG9yLmphdmE=) | 2.81% | [276 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../java/ai/djl/modality/nlp/generate/SeqBatcher.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TZXFCYXRjaGVyLmphdmE=) | 0.75% | [132 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...ity/nlp/generate/ContrastiveSeqBatchScheduler.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9Db250cmFzdGl2ZVNlcUJhdGNoU2NoZWR1bGVyLmphdmE=) | 2.97% | [98 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...i/djl/modality/nlp/generate/SeqBatchScheduler.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TZXFCYXRjaFNjaGVkdWxlci5qYXZh) | 9.83% | [55 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../java/ai/djl/modality/cv/BufferedImageFactory.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0J1ZmZlcmVkSW1hZ2VGYWN0b3J5LmphdmE=) | 40.96% | [47 Missing and 2 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...a/ai/djl/modality/nlp/generate/StepGeneration.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L25scC9nZW5lcmF0ZS9TdGVwR2VuZXJhdGlvbi5qYXZh) | 2.04% | [48 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [api/src/main/java/ai/djl/ndarray/NDArray.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRBcnJheS5qYXZh) | 43.42% | [39 Missing and 4 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...n/java/ai/djl/modality/cv/output/CategoryMask.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9DYXRlZ29yeU1hc2suamF2YQ==) | 22.00% | [39 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [...i/src/main/java/ai/djl/ndarray/NDArrayAdapter.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL25kYXJyYXkvTkRBcnJheUFkYXB0ZXIuamF2YQ==) | 71.21% | [31 Missing and 7 partials :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU2VtYW50aWNTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | 37.50% | [35 Missing :warning: ](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) |
| ... and [73 more](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

:exclamation: Your organization needs to install the [Codecov GitHub app](https://github.com/apps/codecov/installations/select_target) to enable full functionality.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2981      +/-   ##
============================================
+ Coverage     72.08%   72.10%   +0.02%     
- Complexity     5126     7273    +2147     
============================================
  Files           473      722     +249     
  Lines         21970    32517   +10547     
  Branches       2351     3393    +1042     
============================================
+ Hits          15838    23448    +7610     
- Misses         4925     7446    +2521     
- Partials       1207     1623     +416     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/deepjavalibrary/djl/pull/2981?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/.github/workflows/native_jni_s3_pytorch_android.yml b/.github/workflows/native_jni_s3_pytorch_android.yml
index b871f1f39..0376856bb 100644
--- a/.github/workflows/native_jni_s3_pytorch_android.yml
+++ b/.github/workflows/native_jni_s3_pytorch_android.yml
@@ -26,7 +26,7 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-gradle-
       - name: Install NDK
-        run: echo "y" | sudo ${ANDROID_HOME}/tools/bin/sdkmanager --install "ndk;${NDK_VERSION}"
+        run: echo "y" | sudo ${ANDROID_HOME}/cmdline-tools/latest/bin/sdkmanager --install "ndk;${NDK_VERSION}"
       - name: build android
         run: |
           export ANDROID_NDK=${ANDROID_SDK_ROOT}/ndk/${NDK_VERSION}
diff --git a/api/src/main/java/ai/djl/BaseModel.java b/api/src/main/java/ai/djl/BaseModel.java
index 572ab6550..db2f0d3dd 100644
--- a/api/src/main/java/ai/djl/BaseModel.java
+++ b/api/src/main/java/ai/djl/BaseModel.java
@@ -339,8 +339,12 @@ public abstract class BaseModel implements Model {
     protected boolean readParameters(Path paramFile, Map<String, ?> options)
             throws IOException, MalformedModelException {
         logger.debug("Try to load model from {}", paramFile);
-        try (DataInputStream dis =
-                new DataInputStream(new BufferedInputStream(Files.newInputStream(paramFile)))) {
+        return readParameters(Files.newInputStream(paramFile), options);
+    }
+
+    protected boolean readParameters(InputStream paramStream, Map<String, ?> options)
+            throws IOException, MalformedModelException {
+        try (DataInputStream dis = new DataInputStream(new BufferedInputStream(paramStream))) {
             byte[] buf = new byte[4];
             dis.readFully(buf);
             if (!"DJL@".equals(new String(buf, StandardCharsets.US_ASCII))) {
diff --git a/docs/development/inference_performance_optimization.md b/docs/development/inference_performance_optimization.md
index 27bccfd3f..0fdc67d99 100644
--- a/docs/development/inference_performance_optimization.md
+++ b/docs/development/inference_performance_optimization.md
@@ -85,6 +85,23 @@ You can enable it by setting the environment variable:
 
 You might see an exception if a data type or operator is not supported with the oneDNN device.
 
+#### oneDNN(MKLDNN) tuning on AWS Graviton3
+AWS Graviton3(E) (e.g. c7g/m7g/r7g, c7gn and Hpc7g instances) supports BF16 format for ML acceleration. This can be enabled in oneDNN by setting the below environment variable
+```
+grep -q bf16 /proc/cpuinfo && export DNNL_DEFAULT_FPMATH_MODE=BF16
+```
+To avoid redundant primitive creation latency overhead, enable primitive caching by setting the LRU cache capacity. Please note this caching feature increases the memory footprint. It is recommended to tune the capacity to an optimal value for a given use case.
+
+```
+export LRU_CACHE_CAPACITY=1024
+```
+
+In addition to avoiding the redundant allocations, tensor memory allocation latencies can be optimized  with Linux transparent huge pages (THP). To enable THP allocations, set the following torch environment variable.
+```
+export THP_MEM_ALLOC_ENABLE=1
+```
+Please refer to [PyTorch Graviton tutorial](https://pytorch.org/tutorials/recipes/inference_tuning_on_aws_graviton.html) for more details on how to achieve the best PyTorch inference performance on AWS Graviton3 instances.
+
 #### CuDNN acceleration
 PyTorch has a special flag that is used for a CNN or related network speed up. If your input size won't change frequently,
 you may benefit from enabling this configuration in your model:
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtModel.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtModel.java
index 35e95f7de..db0e95e79 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtModel.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtModel.java
@@ -33,6 +33,7 @@ import java.io.InputStream;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.util.ArrayList;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 import java.util.function.Predicate;
@@ -132,7 +133,8 @@ public class PtModel extends BaseModel {
 
     /** {@inheritDoc} */
     @Override
-    public void load(InputStream modelStream, Map<String, ?> options) throws IOException {
+    public void load(InputStream modelStream, Map<String, ?> options)
+            throws IOException, MalformedModelException {
         boolean mapLocation = false;
         if (options != null) {
             mapLocation = Boolean.parseBoolean((String) options.get("mapLocation"));
@@ -146,11 +148,26 @@ public class PtModel extends BaseModel {
      * @param modelStream the stream of the model file
      * @param mapLocation force load to specified device if true
      * @throws IOException model loading error
+     * @throws MalformedModelException if model file is corrupted
      */
-    public void load(InputStream modelStream, boolean mapLocation) throws IOException {
-        modelDir = Files.createTempDirectory("pt-model");
-        modelDir.toFile().deleteOnExit();
-        block = JniUtils.loadModule((PtNDManager) manager, modelStream, mapLocation, false);
+    public void load(InputStream modelStream, boolean mapLocation)
+            throws IOException, MalformedModelException {
+        wasLoaded = true;
+        if (block == null) {
+            modelDir = Files.createTempDirectory("pt-model");
+            modelDir.toFile().deleteOnExit();
+            block = JniUtils.loadModule((PtNDManager) manager, modelStream, mapLocation, false);
+
+            /*
+             * By default, the parameters are frozen, since the previous version before adding this
+             * trainParam, they were frozen due to the setting JITCallGuard guard, which disables
+             * autograd. Also, the pretrained parameters usually should not be updated too much. It
+             * is safe to freeze it. Users may unfreeze it and set their learning rate small.
+             */
+            block.freezeParameters(true);
+        } else {
+            readParameters(modelStream, Collections.emptyMap());
+        }
     }
 
     private Path findModelFile(String... prefixes) {


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/training/ModelTest.java b/integration/src/main/java/ai/djl/integration/tests/training/ModelTest.java
index ca6801290..3ace9c2bd 100644
--- a/integration/src/main/java/ai/djl/integration/tests/training/ModelTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/training/ModelTest.java
@@ -27,6 +27,7 @@ import org.testng.Assert;
 import org.testng.annotations.Test;
 
 import java.io.IOException;
+import java.nio.file.Files;
 import java.nio.file.Paths;
 
 public class ModelTest {
@@ -37,7 +38,9 @@ public class ModelTest {
         block.add(Conv2d.builder().setKernelShape(new Shape(1, 1)).setFilters(10).build());
         block.add(BatchNorm.builder().build());
         try (Model saveModel = Model.newInstance("saveModel", TestUtils.getEngine());
-                Model loadModel = Model.newInstance("loadModel", TestUtils.getEngine())) {
+                Model loadModel = Model.newInstance("loadModel", TestUtils.getEngine());
+                Model loadStreamModel =
+                        Model.newInstance("loadStreamModel", TestUtils.getEngine()); ) {
             block.initialize(saveModel.getNDManager(), DataType.FLOAT32, new Shape(1, 3, 32, 32));
             ParameterList savedParameters = block.getParameters();
             saveModel.setBlock(block);
@@ -48,6 +51,13 @@ public class ModelTest {
             loadModel.load(Paths.get("build/tmp/test/models"), "saveAndLoad");
             ParameterList loadedParameters = loadModel.getBlock().getParameters();
             compareParameters(savedParameters, loadedParameters);
+
+            loadStreamModel.setBlock(block);
+            loadStreamModel.load(
+                    Files.newInputStream(
+                            Paths.get("build/tmp/test/models/saveAndLoad-0000.params")));
+            loadedParameters = loadStreamModel.getBlock().getParameters();
+            compareParameters(savedParameters, loadedParameters);
         }
     }