From f457f66e6e8451c3ababf5857f59a53e65e1fef9 Mon Sep 17 00:00:00 2001 From: Zach Kimberg Date: Wed, 7 Feb 2024 11:44:24 -0800 Subject: [PATCH] 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. --- api/src/main/java/ai/djl/BaseModel.java | 8 ++++-- .../java/ai/djl/pytorch/engine/PtModel.java | 27 +++++++++++++++---- .../integration/tests/training/ModelTest.java | 12 ++++++++- 3 files changed, 39 insertions(+), 8 deletions(-) diff --git a/api/src/main/java/ai/djl/BaseModel.java b/api/src/main/java/ai/djl/BaseModel.java index 572ab65508c..db2f0d3dd70 100644 --- 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 options) throws I protected boolean readParameters(Path paramFile, Map 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 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 index 35e95f7de86..db0e95e79c1 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.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 options) /** {@inheritDoc} */ @Override - public void load(InputStream modelStream, Map options) throws IOException { + public void load(InputStream modelStream, Map 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 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 index ca680129062..3ace9c2bdf5 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.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); } }