Repository: deepjavalibrary/djl
Pull request #1822 (2022-07-21 21:13:41)

[doc] Immutable array output from InferenceMode PyTorch
This doc change originates from [issue 1774](https://github.com/deepjavalibrary/djl/issues/1774). It reminds user that the output array during inference mode is immutable. When they need to change the output NDArray in a implementation of postProcessor, they will know.


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1822?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1822](https://codecov.io/gh/deepjavalibrary/djl/pull/1822?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (75617a8) into [master](https://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) (bb5073f) will **decrease** coverage by `2.48%`.
> The diff coverage is `65.02%`.

```diff
@@             Coverage Diff              @@
##             master    #1822      +/-   ##
============================================
- Coverage     72.08%   69.60%   -2.49%     
- Complexity     5126     5500     +374     
============================================
  Files           473      524      +51     
  Lines         21970    24339    +2369     
  Branches       2351     2652     +301     
============================================
+ Hits          15838    16942    +1104     
- Misses         4925     6087    +1162     
- Partials       1207     1310     +103     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1822?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | Coverage Δ | |
|---|---|---|
| [api/src/main/java/ai/djl/modality/cv/Image.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1822/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L0ltYWdlLmphdmE=) | `69.23% <ø> (-4.11%)` | :arrow_down: |
| [...rc/main/java/ai/djl/modality/cv/MultiBoxPrior.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1822/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L011bHRpQm94UHJpb3IuamF2YQ==) | `76.00% <ø> (ø)` | |
| [...rc/main/java/ai/djl/modality/cv/output/Joints.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1822/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9Kb2ludHMuamF2YQ==) | `71.42% <ø> (ø)` | |
| [.../main/java/ai/djl/modality/cv/output/Landmark.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1822/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9MYW5kbWFyay5qYXZh) | `100.00% <ø> (ø)` | |
| [...main/java/ai/djl/modality/cv/output/Rectangle.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1822/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L291dHB1dC9SZWN0YW5nbGUuamF2YQ==) | `72.41% <0.00%> (ø)` | |
| [...i/djl/modality/cv/translator/BigGANTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1822/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvci5qYXZh) | `21.42% <0.00%> (-5.24%)` | :arrow_down: |
| [...odality/cv/translator/BigGANTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1822/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvQmlnR0FOVHJhbnNsYXRvckZhY3RvcnkuamF2YQ==) | `33.33% <0.00%> (+8.33%)` | :arrow_up: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1822/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `14.28% <0.00%> (-3.90%)` | :arrow_down: |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1822/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU2VtYW50aWNTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../cv/translator/StyleTransferTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/1822/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvU3R5bGVUcmFuc2ZlclRyYW5zbGF0b3JGYWN0b3J5LmphdmE=) | `40.00% <ø> (ø)` | |
| ... and [417 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1822/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1822?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1822?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [654d1cd...75617a8](https://codecov.io/gh/deepjavalibrary/djl/pull/1822?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/api/src/main/java/ai/djl/translate/PostProcessor.java b/api/src/main/java/ai/djl/translate/PostProcessor.java
index ffdb5e6f2..97a7241f4 100644
--- a/api/src/main/java/ai/djl/translate/PostProcessor.java
+++ b/api/src/main/java/ai/djl/translate/PostProcessor.java
@@ -25,7 +25,8 @@ public interface PostProcessor<O> {
      * Processes the output NDList to the corresponding output object.
      *
      * @param ctx the toolkit used for post-processing
-     * @param list the output NDList after inference
+     * @param list the output NDList after inference, usually immutable in engines like
+     *     PyTorch. @see <a href="https://github.com/deepjavalibrary/djl/issues/1774">Issue 1774</a>
      * @return the output object of expected type
      * @throws Exception if an error occurs during processing output
      */
diff --git a/djl-zero/src/main/java/ai/djl/zero/cv/ImageClassification.java b/djl-zero/src/main/java/ai/djl/zero/cv/ImageClassification.java
index 5814fc182..a3548263a 100644
--- a/djl-zero/src/main/java/ai/djl/zero/cv/ImageClassification.java
+++ b/djl-zero/src/main/java/ai/djl/zero/cv/ImageClassification.java
@@ -18,6 +18,7 @@ import ai.djl.Model;
 import ai.djl.basicdataset.cv.classification.ImageClassificationDataset;
 import ai.djl.basicdataset.cv.classification.ImageNet;
 import ai.djl.basicdataset.cv.classification.Mnist;
+import ai.djl.basicmodelzoo.cv.classification.MobileNetV1;
 import ai.djl.basicmodelzoo.cv.classification.ResNetV1;
 import ai.djl.modality.Classifications;
 import ai.djl.modality.cv.Image;
@@ -129,12 +130,21 @@ public final class ImageClassification {
         // Determine the layers based on performance
         int numLayers = performance.switchPerformance(18, 50, 152);
 
-        Block block =
-                ResNetV1.builder()
-                        .setImageShape(imageShape)
-                        .setNumLayers(numLayers)
-                        .setOutSize(classes.size())
-                        .build();
+        Block block;
+
+        if (performance.equals(Performance.FAST)) {
+            // for small and fast cases, build a MobileNet
+            block = MobileNetV1.builder().setOutSize(classes.size()).build();
+        } else {
+            // for large cases, build a ResNet
+            block =
+                    ResNetV1.builder()
+                            .setImageShape(imageShape)
+                            .setNumLayers(numLayers)
+                            .setOutSize(classes.size())
+                            .build();
+        }
+
         Model model = Model.newInstance("ImageClassification");
         model.setBlock(block);
 
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/AlexNet.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/AlexNet.java
index b56267994..d34bf65d1 100644
--- a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/AlexNet.java
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/AlexNet.java
@@ -101,8 +101,8 @@ public final class AlexNet {
                 .add(Activation::relu)
                 .add(Dropout.builder().optRate(builder.dropOutRate).build())
                 // Output layer. The number of
-                // classes is 10, instead of 1000 as in the paper
-                .add(Linear.builder().setUnits(10).build());
+                // classes is 10 as default, instead of 1000 as in the paper
+                .add(Linear.builder().setUnits(builder.outSize).build());
     }
 
     /**
@@ -121,6 +121,8 @@ public final class AlexNet {
         int numLayers = 7;
         int[] numChannels = {96, 256, 384, 384, 256, 4096, 4096};
 
+        long outSize = 10;
+
         Builder() {}
 
         /**
@@ -151,6 +153,17 @@ public final class AlexNet {
             return this;
         }
 
+        /**
+         * Sets the size of the output.
+         *
+         * @param outSize the output size
+         * @return this {@code Builder}
+         */
+        public AlexNet.Builder setOutSize(long outSize) {
+            this.outSize = outSize;
+            return this;
+        }
+
         /**
          * Builds a {@link AlexNet} block.
          *
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/GoogLeNet.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/GoogLeNet.java
index a82083f39..a51ae6fb4 100644
--- a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/GoogLeNet.java
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/GoogLeNet.java
@@ -108,7 +108,7 @@ public final class GoogLeNet {
                         block3,
                         block4,
                         block5,
-                        Linear.builder().setUnits(10).build());
+                        Linear.builder().setUnits(builder.outSize).build());
     }
 
     // c1 - c4 are the number of output channels for each layer in the path
@@ -205,8 +205,21 @@ public final class GoogLeNet {
     /** The Builder to construct a {@link GoogLeNet} object. */
     public static final class Builder {
 
+        long outSize = 10;
+
         Builder() {}
 
+        /**
+         * Sets the size of the output.
+         *
+         * @param outSize the output size
+         * @return this {@code Builder}
+         */
+        public Builder setOutSize(long outSize) {
+            this.outSize = outSize;
+            return this;
+        }
+
         /**
          * Builds a {@link GoogLeNet} block.
          *
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/LeNet.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/LeNet.java
index 7fc6ca410..9128a43ba 100644
--- a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/LeNet.java
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/LeNet.java
@@ -67,7 +67,7 @@ public final class LeNet {
                 .add(Activation::sigmoid)
                 .add(Linear.builder().setUnits(builder.numChannels[3]).build())
                 .add(Activation::sigmoid)
-                .add(Linear.builder().setUnits(10).build());
+                .add(Linear.builder().setUnits(builder.outSize).build());
     }
 
     /**
@@ -84,6 +84,7 @@ public final class LeNet {
 
         int numLayers = 4;
         int[] numChannels = {6, 16, 120, 84};
+        long outSize = 10;
 
         Builder() {}
 
@@ -104,6 +105,17 @@ public final class LeNet {
             return this;
         }
 
+        /**
+         * Sets the size of the output.
+         *
+         * @param outSize the output size
+         * @return this {@code Builder}
+         */
+        public Builder setOutSize(long outSize) {
+            this.outSize = outSize;
+            return this;
+        }
+
         /**
          * Builds a {@link LeNet} block.
          *
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/MobileNetV1.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/MobileNetV1.java
new file mode 100644
index 000000000..7477f65c2
--- /dev/null
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/MobileNetV1.java
@@ -0,0 +1,275 @@
+/*
+ * Copyright 2022 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.cv.classification;
+
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Activation;
+import ai.djl.nn.Block;
+import ai.djl.nn.SequentialBlock;
+import ai.djl.nn.convolutional.Conv2d;
+import ai.djl.nn.core.Linear;
+import ai.djl.nn.norm.BatchNorm;
+import ai.djl.nn.pooling.Pool;
+
+/**
+ * {@code MobileNetV1} contains a generic implementation of Mobilenet adapted from
+ * https://github.com/weiaicunzai/pytorch-cifar100/blob/master/models/mobilenet.py (Original author
+ * weiaicunzai).
+ *
+ * <p>see https://arxiv.org/pdf/1704.04861.pdf for more information about MobileNet
+ */
+public final class MobileNetV1 {
+    static final int[] FILTERS = {32, 64, 128, 128, 256, 256, 512, 512, 1024, 1024};
+
+    private MobileNetV1() {}
+
+    /**
+     * Builds a {@link Block} that represent a depthWise-pointWise Unit used in the implementation
+     * of the MobileNet Model.
+     *
+     * @param inputChannels number of inputChannels, used for depthWise Kernel
+     * @param outputChannels number of outputChannels, used for pointWise kernel
+     * @param stride control the stride of depthWise Kernel
+     * @param builder add the builder to obtain batchNormMomentum
+     * @return a {@link Block} that represent a depthWise-pointWise Unit
+     */
+    public static Block depthSeparableConv2d(
+            int inputChannels, int outputChannels, int stride, Builder builder) {
+        // depthWise does not include bias
+        SequentialBlock depthWise = new SequentialBlock();
+        depthWise
+                .add(
+                        Conv2d.builder()
+                                .setKernelShape(
+                                        new Shape(3, 3)) // the kernel size of depthWise is always 3
+                                .optBias(false)
+                                .optPadding(new Shape(1, 1)) // padding = same
+                                .optStride(new Shape(stride, stride)) // stride is either 2 or 1
+                                .optGroups(
+                                        inputChannels) // depthWise with 1 filter per input channel
+                                .setFilters(inputChannels)
+                                .build())
+                .add( // add a batchNorm
+                        BatchNorm.builder()
+                                .optEpsilon(2E-5f)
+                                .optMomentum(builder.batchNormMomentum)
+                                .build())
+                .add(Activation.reluBlock());
+
+        SequentialBlock pointWise = new SequentialBlock();
+        pointWise
+                .add(
+                        Conv2d.builder()
+                                .setKernelShape(new Shape(1, 1)) // no padding or stride
+                                .setFilters(outputChannels)
+                                .optBias(false)
+                                .build())
+                .add(
+                        BatchNorm.builder()
+                                .optEpsilon(2E-5f)
+                                .optMomentum(builder.batchNormMomentum)
+                                .build())
+                .add(Activation.reluBlock());
+
+        return depthWise.add(pointWise); // two blocks are merged together
+    }
+
+    /**
+     * Creates a new {@link Block} of {@link MobileNetV1} with the arguments from the given {@link
+     * Builder}.
+     *
+     * @param builder the {@link Builder} with the necessary arguments
+     * @return a {@link Block} that represents the required MobileNet model
+     */
+    public static Block mobilenet(Builder builder) {
+        // no bias in MobileNet
+        SequentialBlock mobileNet = new SequentialBlock();
+
+        mobileNet
+                .add(
+                        // conv1
+                        new SequentialBlock()
+                                .add(
+                                        Conv2d.builder()
+                                                .setKernelShape(new Shape(3, 3))
+                                                .optBias(false)
+                                                .optStride(new Shape(2, 2))
+                                                .optPadding(new Shape(1, 1)) // padding = 'same'
+                                                .setFilters(
+                                                        (int)
+                                                                (FILTERS[0]
+                                                                        * builder.widthMultiplier))
+                                                .build())
+                                .add(
+                                        BatchNorm.builder()
+                                                .optEpsilon(2E-5f)
+                                                .optMomentum(builder.batchNormMomentum)
+                                                .build())
+                                .add(Activation.reluBlock()))
+                // separable conv1
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[0] * builder.widthMultiplier),
+                                (int) (FILTERS[1] * builder.widthMultiplier),
+                                1,
+                                builder))
+                // separable conv2
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[1] * builder.widthMultiplier),
+                                (int) (FILTERS[2] * builder.widthMultiplier),
+                                2,
+                                builder))
+                // separable conv3
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[2] * builder.widthMultiplier),
+                                (int) (FILTERS[3] * builder.widthMultiplier),
+                                1,
+                                builder))
+                // separable conv4
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[3] * builder.widthMultiplier),
+                                (int) (FILTERS[4] * builder.widthMultiplier),
+                                2,
+                                builder))
+                // separable conv5
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[4] * builder.widthMultiplier),
+                                (int) (FILTERS[5] * builder.widthMultiplier),
+                                1,
+                                builder))
+                // separable conv6
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[5] * builder.widthMultiplier),
+                                (int) (FILTERS[6] * builder.widthMultiplier),
+                                2,
+                                builder))
+                // separable conv7*5
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[6] * builder.widthMultiplier),
+                                (int) (FILTERS[7] * builder.widthMultiplier),
+                                1,
+                                builder))
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[6] * builder.widthMultiplier),
+                                (int) (FILTERS[7] * builder.widthMultiplier),
+                                1,
+                                builder))
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[6] * builder.widthMultiplier),
+                                (int) (FILTERS[7] * builder.widthMultiplier),
+                                1,
+                                builder))
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[6] * builder.widthMultiplier),
+                                (int) (FILTERS[7] * builder.widthMultiplier),
+                                1,
+                                builder))
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[6] * builder.widthMultiplier),
+                                (int) (FILTERS[7] * builder.widthMultiplier),
+                                1,
+                                builder))
+                // separable conv8
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[7] * builder.widthMultiplier),
+                                (int) (FILTERS[8] * builder.widthMultiplier),
+                                2,
+                                builder))
+                // separable conv9
+                .add(
+                        depthSeparableConv2d(
+                                (int) (FILTERS[8] * builder.widthMultiplier),
+                                (int) (FILTERS[9] * builder.widthMultiplier),
+                                1,
+                                builder)) // maybe the paper goes wrong here
+                // AveragePool
+                .add(Pool.globalAvgPool2dBlock())
+                // FC
+                .add(Linear.builder().setUnits(builder.outSize).build());
+
+        return mobileNet;
+    }
+
+    /**
+     * Creates a builder to build a {@link MobileNetV1}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new MobileNetV1.Builder();
+    }
+
+    /** The Builder to construct a {@link MobileNetV1} object. */
+    public static final class Builder {
+
+        float batchNormMomentum = 0.9f;
+        float widthMultiplier = 1f; // width multiplier(also named alpha) defined in the paper
+        long outSize = 10; // 10 as default for basic datasets like cifar-10 or mnist
+
+        Builder() {}
+
+        /**
+         * Sets the widthMultiplier of MobileNet.
+         *
+         * @param widthMultiplier the widthMultiplier of MobileNet
+         * @return this {@code Builder}
+         */
+        public Builder optWidthMultiplier(float widthMultiplier) {
+            this.widthMultiplier = widthMultiplier;
+            return this;
+        }
+
+        /**
+         * Sets the momentum of batchNorm layer.
+         *
+         * @param batchNormMomentum the momentum
+         * @return this {@code Builder}
+         */
+        public Builder optBatchNormMomentum(float batchNormMomentum) {
+            this.batchNormMomentum = batchNormMomentum;
+            return this;
+        }
+
+        /**
+         * Sets the size of the output.
+         *
+         * @param outSize the output size
+         * @return this {@code Builder}
+         */
+        public Builder setOutSize(long outSize) {
+            this.outSize = outSize;
+            return this;
+        }
+
+        /**
+         * Builds a {@link MobileNetV1} block.
+         *
+         * @return the {@link MobileNetV1} block
+         */
+        public Block build() {
+            return mobilenet(this);
+        }
+    }
+}
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/VGG.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/VGG.java
index be6c1c660..ca33e57a1 100644
--- a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/VGG.java
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/VGG.java
@@ -57,7 +57,7 @@ public final class VGG {
                 .add(Linear.builder().setUnits(4096).build())
                 .add(Activation::relu)
                 .add(Dropout.builder().optRate(0.5f).build())
-                .add(Linear.builder().setUnits(10).build());
+                .add(Linear.builder().setUnits(builder.outSize).build());
 
         return block;
     }
@@ -101,6 +101,7 @@ public final class VGG {
 
         int numLayers = 11;
         int[][] convArch = {{1, 64}, {1, 128}, {2, 256}, {2, 512}, {2, 512}};
+        long outSize = 10;
 
         /**
          * Sets the number of layers. It is equal to total sum of numConvs in convArch + 3.
@@ -137,6 +138,17 @@ public final class VGG {
             return this;
         }
 
+        /**
+         * Sets the size of the output.
+         *
+         * @param outSize the output size
+         * @return this {@code Builder}
+         */
+        public Builder setOutSize(long outSize) {
+            this.outSize = outSize;
+            return this;
+        }
+
         /**
          * Builds a {@link VGG} block.
          *


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/MobileNetV1Test.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/MobileNetV1Test.java
new file mode 100644
index 000000000..0ab2f2a7c
--- /dev/null
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/MobileNetV1Test.java
@@ -0,0 +1,95 @@
+/*
+ * Copyright 2022 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.integration.tests.model_zoo.classification;
+
+import ai.djl.Model;
+import ai.djl.basicmodelzoo.cv.classification.MobileNetV1;
+import ai.djl.engine.Engine;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Block;
+import ai.djl.nn.Parameter;
+import ai.djl.testing.Assertions;
+import ai.djl.training.DefaultTrainingConfig;
+import ai.djl.training.EasyTrain;
+import ai.djl.training.Trainer;
+import ai.djl.training.TrainingConfig;
+import ai.djl.training.dataset.Batch;
+import ai.djl.training.initializer.Initializer;
+import ai.djl.training.loss.Loss;
+import ai.djl.translate.Batchifier;
+import ai.djl.util.PairList;
+
+import org.testng.annotations.Test;
+
+public class MobileNetV1Test {
+    @Test
+    public void testTrain() {
+        TrainingConfig config =
+                new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
+                        .optDevices(Engine.getInstance().getDevices(2))
+                        .optInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+        Block mobilenet = MobileNetV1.builder().build();
+        try (Model model = Model.newInstance("mobilenet")) {
+            model.setBlock(mobilenet);
+            try (Trainer trainer = model.newTrainer(config)) {
+                int batchSize = 1;
+                Shape inputShape = new Shape(batchSize, 3, 224, 224);
+                trainer.initialize(inputShape);
+                NDManager manager = trainer.getManager();
+                NDArray input = manager.ones(inputShape);
+                NDArray label = manager.ones(new Shape(batchSize, 1));
+                Batch batch =
+                        new Batch(
+                                manager.newSubManager(),
+                                new NDList(input),
+                                new NDList(label),
+                                batchSize,
+                                Batchifier.STACK,
+                                Batchifier.STACK,
+                                0,
+                                0);
+                PairList<String, Parameter> parameters = mobilenet.getParameters();
+                EasyTrain.trainBatch(trainer, batch);
+                trainer.step();
+
+                // expected kernelShapes
+                NDArray expectedAtIndex0 =
+                        manager.ones(new Shape(32, 3, 3, 3)); // 32*3*3*3 for first layer
+                NDArray expectedAtIndex5 =
+                        manager.ones(new Shape(32, 1, 3, 3)); // 32*1*3*3  for depthWiseLayer1
+                NDArray expectedAtIndex120 =
+                        manager.ones(
+                                new Shape(
+                                        1024, 512, 1,
+                                        1)); // 1024*512*1*1  for pointWiseLayer at last
+                NDArray expectedAvgPool =
+                        manager.ones(
+                                new Shape(10, 1024)); // outSize = 10, AvgPoolSize should be 10*1024
+
+                // test if the shape of the kernel sizes are right
+                Assertions.assertAlmostEquals(
+                        parameters.get(0).getValue().getArray(), expectedAtIndex0);
+                Assertions.assertAlmostEquals(
+                        parameters.get(5).getValue().getArray(), expectedAtIndex5);
+                Assertions.assertAlmostEquals(
+                        parameters.get(120).getValue().getArray(), expectedAtIndex120);
+                Assertions.assertAlmostEquals(
+                        parameters.get(135).getValue().getArray(), expectedAvgPool);
+            }
+        }
+    }
+}