Refactor BlockFactory interface

api/src/main/java/ai/djl/BaseModel.java

@@ -19,13 +19,11 @@
 import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.nn.Block;
-import ai.djl.nn.BlockFactory;
 import ai.djl.nn.SymbolBlock;
 import ai.djl.training.ParameterStore;
 import ai.djl.training.Trainer;
 import ai.djl.training.TrainingConfig;
 import ai.djl.translate.Translator;
-import ai.djl.util.ClassLoaderUtils;
 import ai.djl.util.Pair;
 import ai.djl.util.PairList;
 import ai.djl.util.Utils;
@@ -217,14 +215,6 @@ protected void setModelDir(Path modelDir) {
         this.modelDir = modelDir.toAbsolutePath();
     }
 
-    protected Block loadFromBlockFactory() {
-        BlockFactory factory = ClassLoaderUtils.findImplementation(modelDir, null);
-        if (factory == null) {
-            return null;
-        }
-        return factory.newBlock(manager);
-    }
-
     /** {@inheritDoc} */
     @Override
     public void save(Path modelPath, String newModelName) throws IOException {

api/src/main/java/ai/djl/nn/BlockFactory.java

@@ -12,9 +12,12 @@
  */
 package ai.djl.nn;
 
-import ai.djl.ndarray.NDManager;
+import ai.djl.Model;
 import ai.djl.repository.zoo.ModelZoo;
+import java.io.IOException;
 import java.io.Serializable;
+import java.nio.file.Path;
+import java.util.Map;
 
 /**
  * Block factory is a component to make standard for block creating and saving procedure. Block
@@ -27,8 +30,11 @@ public interface BlockFactory extends Serializable {
     /**
      * Constructs the uninitialized block.
      *
-     * @param manager the manager to assign to block
+     * @param model the model of the block
+     * @param modelPath the directory of the model location
+     * @param arguments the block creation arguments
      * @return the uninitialized block
+     * @throws IOException if IO operation fails during creating block
      */
-    Block newBlock(NDManager manager);
+    Block newBlock(Model model, Path modelPath, Map<String, ?> arguments) throws IOException;
 }

api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java

@@ -23,6 +23,8 @@
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.translator.ImageClassificationTranslatorFactory;
 import ai.djl.ndarray.NDList;
+import ai.djl.nn.Block;
+import ai.djl.nn.BlockFactory;
 import ai.djl.repository.Artifact;
 import ai.djl.repository.MRL;
 import ai.djl.repository.Repository;
@@ -32,6 +34,7 @@
 import ai.djl.translate.TranslateException;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
+import ai.djl.util.ClassLoaderUtils;
 import ai.djl.util.Pair;
 import ai.djl.util.Progress;
 import java.io.IOException;
@@ -155,10 +158,14 @@ public <I, O> ZooModel<I, O> loadModel(Criteria<I, O> criteria)
                 modelName = artifact.getName();
             }
 
-            Model model = createModel(modelName, criteria.getDevice(), artifact, arguments, engine);
-            if (criteria.getBlock() != null) {
-                model.setBlock(criteria.getBlock());
-            }
+            Model model =
+                    createModel(
+                            modelPath,
+                            modelName,
+                            criteria.getDevice(),
+                            criteria.getBlock(),
+                            arguments,
+                            engine);
             model.load(modelPath, null, options);
             Translator<I, O> translator = factory.newInstance(model, arguments);
             return new ZooModel<>(model, translator);
@@ -182,13 +189,25 @@ public List<Artifact> listModels() throws IOException {
     }
 
     protected Model createModel(
+            Path modelPath,
             String name,
             Device device,
-            Artifact artifact,
+            Block block,
             Map<String, Object> arguments,
             String engine)
             throws IOException {
-        return Model.newInstance(name, device, engine);
+        Model model = Model.newInstance(name, device, engine);
+        if (block == null) {
+            String className = (String) arguments.get("blockFactory");
+            BlockFactory factory = ClassLoaderUtils.findImplementation(modelPath, className);
+            if (factory != null) {
+                block = factory.newBlock(model, modelPath, arguments);
+            }
+        }
+        if (block != null) {
+            model.setBlock(block);
+        }
+        return model;
     }
 
     /** {@inheritDoc} */

extensions/fasttext/src/main/java/ai/djl/fasttext/zoo/nlp/textclassification/TextClassificationModelLoader.java

@@ -19,7 +19,7 @@
 import ai.djl.fasttext.FtModel;
 import ai.djl.fasttext.zoo.FtModelZoo;
 import ai.djl.modality.Classifications;
-import ai.djl.repository.Artifact;
+import ai.djl.nn.Block;
 import ai.djl.repository.MRL;
 import ai.djl.repository.Repository;
 import ai.djl.repository.zoo.BaseModelLoader;
@@ -30,6 +30,7 @@
 import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.Pair;
 import java.io.IOException;
+import java.nio.file.Path;
 import java.util.Map;
 
 /** Model loader for fastText cooking stackexchange models. */
@@ -68,9 +69,10 @@ public ZooModel<String, Classifications> loadModel()
     /** {@inheritDoc} */
     @Override
     protected Model createModel(
+            Path modelPath,
             String name,
             Device device,
-            Artifact artifact,
+            Block block,
             Map<String, Object> arguments,
             String engine) {
         return new FtModel(name);

integration/src/main/java/ai/djl/integration/tests/nn/BlockFactoryTest.java

@@ -30,59 +30,21 @@
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ModelNotFoundException;
 import ai.djl.repository.zoo.ZooModel;
-import ai.djl.testing.Assertions;
 import ai.djl.training.ParameterStore;
 import ai.djl.training.util.ProgressBar;
-import ai.djl.translate.NoopTranslator;
 import ai.djl.translate.TranslateException;
 import ai.djl.util.Utils;
 import ai.djl.util.ZipUtils;
-import java.io.ByteArrayInputStream;
-import java.io.ByteArrayOutputStream;
-import java.io.DataInputStream;
-import java.io.DataOutputStream;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Map;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
 public class BlockFactoryTest {
 
-    @Test
-    public void testBlockLoadingSaving()
-            throws IOException, ModelNotFoundException, MalformedModelException,
-                    TranslateException {
-        TestBlockFactory factory = new TestBlockFactory();
-        Model model = factory.getRemoveLastBlockModel();
-        try (NDManager manager = NDManager.newBaseManager()) {
-            Block block = model.getBlock();
-            block.forward(
-                    new ParameterStore(manager, true),
-                    new NDList(manager.ones(new Shape(1, 3, 32, 32))),
-                    true);
-            ByteArrayOutputStream os = new ByteArrayOutputStream();
-            block.saveParameters(new DataOutputStream(os));
-            ByteArrayInputStream bis = new ByteArrayInputStream(os.toByteArray());
-            Block newBlock = factory.newBlock(manager);
-            newBlock.loadParameters(manager, new DataInputStream(bis));
-            try (Model test = Model.newInstance("test")) {
-                test.setBlock(newBlock);
-                try (Predictor<NDList, NDList> predOrigin =
-                                model.newPredictor(new NoopTranslator());
-                        Predictor<NDList, NDList> predDest =
-                                test.newPredictor(new NoopTranslator())) {
-                    NDList input = new NDList(manager.ones(new Shape(1, 3, 32, 32)));
-                    NDList originOut = predOrigin.predict(input);
-                    NDList destOut = predDest.predict(input);
-                    Assertions.assertAlmostEquals(originOut, destOut);
-                }
-            }
-        }
-        model.close();
-    }
-
     @Test
     public void testBlockFactoryLoadingFromZip()
             throws MalformedModelException, ModelNotFoundException, IOException,
@@ -97,9 +59,9 @@ public void testBlockFactoryLoadingFromZip()
                         .optModelPath(zipPath)
                         .optModelName("exported")
                         .build();
-        try (NDManager manager = NDManager.newBaseManager();
-                ZooModel<NDList, NDList> model = criteria.loadModel();
+        try (ZooModel<NDList, NDList> model = criteria.loadModel();
                 Predictor<NDList, NDList> pred = model.newPredictor()) {
+            NDManager manager = model.getNDManager();
             NDList destOut = pred.predict(new NDList(manager.ones(new Shape(1, 3, 32, 32))));
             Assert.assertEquals(destOut.singletonOrThrow().getShape(), new Shape(1, 10));
         }
@@ -136,9 +98,9 @@ public static class TestBlockFactory implements BlockFactory {
         private static final long serialVersionUID = 1234567L;
 
         @Override
-        public Block newBlock(NDManager manager) {
+        public Block newBlock(Model model, Path modelPath, Map<String, ?> arguments) {
             SequentialBlock newBlock = new SequentialBlock();
-            newBlock.add(SymbolBlock.newInstance(manager));
+            newBlock.add(SymbolBlock.newInstance(model.getNDManager()));
             newBlock.add(Linear.builder().setUnits(10).build());
             return newBlock;
         }

model-zoo/src/main/java/ai/djl/basicmodelzoo/BasicModelZoo.java

@@ -12,10 +12,10 @@
  */
 package ai.djl.basicmodelzoo;
 
-import ai.djl.basicmodelzoo.cv.classification.MlpModelLoader;
-import ai.djl.basicmodelzoo.cv.classification.ResNetModelLoader;
-import ai.djl.basicmodelzoo.cv.object_detection.ssd.SsdModelLoader;
+import ai.djl.modality.cv.zoo.ImageClassificationModelLoader;
+import ai.djl.modality.cv.zoo.ObjectDetectionModelLoader;
 import ai.djl.repository.Repository;
+import ai.djl.repository.zoo.ModelLoader;
 import ai.djl.repository.zoo.ModelZoo;
 import java.util.HashSet;
 import java.util.Set;
@@ -25,11 +25,15 @@ public class BasicModelZoo implements ModelZoo {
 
     private static final String REPO_URL = "https://mlrepo.djl.ai/";
     private static final Repository REPOSITORY = Repository.newInstance("zoo", REPO_URL);
+    private static final ModelZoo ZOO = new BasicModelZoo();
     public static final String GROUP_ID = "ai.djl.zoo";
 
-    public static final ResNetModelLoader RESNET = new ResNetModelLoader(REPOSITORY);
-    public static final MlpModelLoader MLP = new MlpModelLoader(REPOSITORY);
-    public static final SsdModelLoader SSD = new SsdModelLoader(REPOSITORY);
+    public static final ModelLoader RESNET =
+            new ImageClassificationModelLoader(REPOSITORY, GROUP_ID, "resnet", "0.0.2", ZOO);
+    public static final ModelLoader MLP =
+            new ImageClassificationModelLoader(REPOSITORY, GROUP_ID, "mlp", "0.0.3", ZOO);
+    public static final ModelLoader SSD =
+            new ObjectDetectionModelLoader(REPOSITORY, GROUP_ID, "ssd", "0.0.2", ZOO);
 
     /** {@inheritDoc} */
     @Override

model-zoo/src/main/java/ai/djl/basicmodelzoo/basic/MlpBlockFactory.java

@@ -0,0 +1,49 @@
+/*
+ * Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.basicmodelzoo.basic;
+
+import ai.djl.Model;
+import ai.djl.nn.Block;
+import ai.djl.nn.BlockFactory;
+import java.nio.file.Path;
+import java.util.List;
+import java.util.Map;
+
+/** A {@link BlockFactory} class that creates MLP block. */
+public class MlpBlockFactory implements BlockFactory {
+
+    private static final long serialVersionUID = 1L;
+
+    /** {@inheritDoc} */
+    @Override
+    @SuppressWarnings("unchecked")
+    public Block newBlock(Model model, Path modelPath, Map<String, ?> arguments) {
+        Double width = (Double) arguments.get("width");
+        if (width == null) {
+            width = 28d;
+        }
+        Double height = (Double) arguments.get("height");
+        if (height == null) {
+            height = 28d;
+        }
+        int input = width.intValue() * height.intValue();
+        int output = ((Double) arguments.get("output")).intValue();
+        int[] hidden =
+                ((List<Double>) arguments.get("hidden"))
+                        .stream()
+                        .mapToInt(Double::intValue)
+                        .toArray();
+
+        return new Mlp(input, output, hidden);
+    }
+}