Supports PyTorch stream imperative model load (#2981)

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.
deepjavalibrary · Feb 7, 2024 · f457f66 · f457f66
1 parent 91a69ed
commit f457f66
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 8 deletions.
diff --git a/api/src/main/java/ai/djl/BaseModel.java b/api/src/main/java/ai/djl/BaseModel.java
@@ -339,8 +339,12 @@ protected Path paramPathResolver(String prefix, Map<String, ?> options) throws I
     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/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.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 void load(Path modelPath, String prefix, Map<String, ?> options)
 
     /** {@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 void load(InputStream modelStream, Map<String, ?> options) throws IOExcep
      * @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) {

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
@@ -27,6 +27,7 @@
 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 void testModelSaveAndLoad() throws IOException, MalformedModelException {
         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 void testModelSaveAndLoad() throws IOException, MalformedModelException {
             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);
         }
     }