Remove performance issues from freezing MXNet (#2394)

* Remove performance issues from freezing MXNet In #2360, the behavior of using pre-trained models was to freeze parameters. However freezing the parameters on MXNet seems to cause a significant performance regression for training. This removes those changes for a temporary workaround until a deeper investigation can take place. Co-authored-by: Frank Liu <frankfliu2000@gmail.com>
deepjavalibrary · Feb 14, 2023 · 04cf346 · 04cf346
1 parent 7ef827a
commit 04cf346
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxModel.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxModel.java
@@ -87,6 +87,7 @@ public class MxModel extends BaseModel {
      * @throws IOException Exception for file loading
      */
     @Override
+    @SuppressWarnings("PMD.EmptyControlStatement")
     public void load(Path modelPath, String prefix, Map<String, ?> options)
             throws IOException, MalformedModelException {
         setModelDir(modelPath);
@@ -143,12 +144,15 @@ public void load(Path modelPath, String prefix, Map<String, ?> options)
         boolean trainParam =
                 options != null && Boolean.parseBoolean((String) options.get("trainParam"));
         if (!trainParam) {
-            block.freezeParameters(true);
+            // TODO: See https://github.com/deepjavalibrary/djl/pull/2360
+            // NOPMD
+            // block.freezeParameters(true);
         }
     }
 
     /** {@inheritDoc} */
     @Override
+    @SuppressWarnings("PMD.EmptyControlStatement")
     public Trainer newTrainer(TrainingConfig trainingConfig) {
         PairList<Initializer, Predicate<Parameter>> initializer = trainingConfig.getInitializers();
         if (block == null) {
@@ -157,7 +161,8 @@ public Trainer newTrainer(TrainingConfig trainingConfig) {
         }
         if (wasLoaded) {
             // Unfreeze parameters if training directly
-            block.freezeParameters(false);
+            // TODO: See https://github.com/deepjavalibrary/djl/pull/2360
+            // block.freezeParameters(false);
         }
         for (Pair<Initializer, Predicate<Parameter>> pair : initializer) {
             if (pair.getKey() != null && pair.getValue() != null) {

diff --git a/examples/src/main/java/ai/djl/examples/training/transferlearning/TrainResnetWithCifar10.java b/examples/src/main/java/ai/djl/examples/training/transferlearning/TrainResnetWithCifar10.java
@@ -145,7 +145,6 @@ private static Model getModel(Arguments arguments)
             SequentialBlock newBlock = new SequentialBlock();
             SymbolBlock block = (SymbolBlock) model.getBlock();
             block.removeLastBlock();
-            block.freezeParameters(false);
             newBlock.add(block);
             // the original model don't include the flatten
             // so apply the flatten here