Repository: deepjavalibrary/djl
Pull request #1847 (2022-07-31 13:09:54)

Add MobileNetV2 into ModelZoo and corresponding Test Class MobileNetV…
Add a new class MobileNetV2 into ModelZoo and its corresponding test class MobileNetV2Test into intergration. I think it's almost complete except for some minor problems.It seems like my Block MobileNetV2 will fail when it comes to trainer.step(). The following code is when it goes wrong.

                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();


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/1847?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
> Merging [#1847](https://codecov.io/gh/deepjavalibrary/djl/pull/1847?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) (d23cad5) 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.38%`.
> The diff coverage is `64.61%`.

```diff
@@             Coverage Diff              @@
##             master    #1847      +/-   ##
============================================
- Coverage     72.08%   69.70%   -2.39%     
- Complexity     5126     5525     +399     
============================================
  Files           473      526      +53     
  Lines         21970    24446    +2476     
  Branches       2351     2665     +314     
============================================
+ Hits          15838    17041    +1203     
- Misses         4925     6093    +1168     
- Partials       1207     1312     +105     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/1847?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/1847/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/1847/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/1847/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/1847/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/1847/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/1847/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/1847/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/1847/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/1847/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/1847/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 [420 more](https://codecov.io/gh/deepjavalibrary/djl/pull/1847/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/1847?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/1847?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Last update [17f80f8...d23cad5](https://codecov.io/gh/deepjavalibrary/djl/pull/1847?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).

zachgk: Also, is this performing better then MobileNetV1 during tests? It should be because V2 is an improvement on V1. Then, you should replace the V1 in DJL Zero
zachgk: It looks good to merge. Can you rebase it onto master?

Patch:
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/MobileNetV2.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/MobileNetV2.java
new file mode 100644
index 000000000..1a546033f
--- /dev/null
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/classification/MobileNetV2.java
@@ -0,0 +1,336 @@
+/*
+ * 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.NDArrays;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Activation;
+import ai.djl.nn.Block;
+import ai.djl.nn.Blocks;
+import ai.djl.nn.ParallelBlock;
+import ai.djl.nn.SequentialBlock;
+import ai.djl.nn.convolutional.Conv2d;
+import ai.djl.nn.norm.BatchNorm;
+import ai.djl.nn.pooling.Pool;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * {@code MobileNetV2} contains a generic implementation of MobilenetV2 adapted from
+ * https://github.com/weiaicunzai/pytorch-cifar100/blob/master/models/mobilenetv2.py (Original
+ * author weiaicunzai).
+ *
+ * <p>see https://arxiv.org/pdf/1801.04381.pdf for more information about MobileNetV2
+ */
+public final class MobileNetV2 {
+
+    public static final int FILTERLENGTH = 9;
+    public static final int REPEATLENGTH = 9;
+    public static final int STRIDELENGTH = 9;
+    public static final int MULTILENGTH = 7;
+
+    private MobileNetV2() {}
+    /**
+     * Builds a {@link Block} that represent an inverted residual Unit used in the implementation of
+     * the MobileNetV2 Model.
+     *
+     * @param inputChannels number of inputChannels of the block
+     * @param outputChannels number of outputChannels of the block
+     * @param stride control the stride of a depthWise kernel
+     * @param t the multiTime of the first pointWise Block
+     * @param batchNormMomentum the momentum of batchNormLayer
+     * @return a {@link Block} that represent an inverted residual Unit
+     */
+    public static Block linearBottleNeck(
+            int inputChannels, int outputChannels, int stride, int t, float batchNormMomentum) {
+        SequentialBlock block = new SequentialBlock();
+        block.add(
+                        Conv2d.builder() // PointWise
+                                .setFilters(inputChannels * t)
+                                .setKernelShape(new Shape(1, 1))
+                                .optBias(false)
+                                .build())
+                // add a batchNorm
+                .add(BatchNorm.builder().optEpsilon(2E-5f).optMomentum(batchNormMomentum).build())
+                // add a relu
+                .add(Activation.reluBlock())
+                .add(
+                        Conv2d.builder() // DepthWise
+                                .setKernelShape(new Shape(3, 3))
+                                .setFilters(inputChannels * t)
+                                .optStride(new Shape(stride, stride))
+                                .optPadding(new Shape(1, 1))
+                                .optGroups(inputChannels * t)
+                                .optBias(false)
+                                .build())
+                .add(BatchNorm.builder().optEpsilon(2E-5f).optMomentum(batchNormMomentum).build())
+                .add(Activation.reluBlock())
+                .add(
+                        Conv2d.builder() // PointWise
+                                .setFilters(outputChannels)
+                                .setKernelShape(new Shape(1, 1))
+                                .optBias(false)
+                                .build())
+                .add(BatchNorm.builder().optEpsilon(2E-5f).optMomentum(batchNormMomentum).build());
+
+        // if dimMatch,then add x
+        if (stride == 1 && inputChannels == outputChannels) {
+            return new ParallelBlock(
+                    list ->
+                            new NDList(
+                                    NDArrays.add(
+                                            list.get(0).singletonOrThrow(),
+                                            list.get(1).singletonOrThrow())),
+                    Arrays.asList(block, Blocks.identityBlock()));
+        }
+        return block; // No reluBlock here
+    }
+
+    /**
+     * Builds a {@link Block} that represent multiple repeats of an inverted residual Unit.
+     *
+     * @param repeat the repeatTimes of an inverted residual Block
+     * @param inputChannels number of inputChannels of the block
+     * @param outputChannels number of outputChannels of the block
+     * @param stride the stride of an inverted residual Unit
+     * @param t the multipleTime of a pointWise Kernel
+     * @param batchNormMomentum the momentum of batchNormLayer
+     * @return a {@link Block} that represent several repeated inverted residual Units
+     */
+    public static Block makeStage(
+            int repeat,
+            int inputChannels,
+            int outputChannels,
+            int stride,
+            int t,
+            float batchNormMomentum) {
+        SequentialBlock layers = new SequentialBlock();
+        layers.add(linearBottleNeck(inputChannels, outputChannels, stride, t, batchNormMomentum));
+        for (int i = 0; i < repeat - 1; i++) {
+            layers.add(linearBottleNeck(outputChannels, outputChannels, 1, t, batchNormMomentum));
+        }
+        return layers;
+    }
+
+    /**
+     * Creates a new {@link Block} of {@link MobileNetV2} 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 MobileNetV2 model
+     */
+    public static Block mobilenetV2(Builder builder) {
+        SequentialBlock mobileNet = new SequentialBlock();
+        SequentialBlock pre = new SequentialBlock();
+        for (int i = 0; i < builder.repeatTimes[0]; i++) { // add as a sequence
+            pre.add(
+                            Conv2d.builder()
+                                    .setKernelShape(new Shape(1, 1))
+                                    .setFilters(builder.filters[0])
+                                    .optStride(new Shape(builder.strides[0], builder.strides[0]))
+                                    .optBias(false)
+                                    .build())
+                    .add(
+                            BatchNorm.builder()
+                                    .optEpsilon(2E-5f)
+                                    .optMomentum(builder.batchNormMomentum)
+                                    .build())
+                    .add(Activation.reluBlock());
+        }
+
+        List<Block> bottleNecks = new ArrayList<>();
+        for (int i = 0; i < MULTILENGTH; i++) {
+            bottleNecks.add(
+                    makeStage(
+                            builder.repeatTimes[i + 1],
+                            builder.filters[i],
+                            builder.filters[i + 1],
+                            builder.strides[i + 1],
+                            builder.multiTimes[i],
+                            builder.batchNormMomentum));
+        }
+
+        SequentialBlock conv1 = new SequentialBlock();
+        for (int i = 0; i < builder.repeatTimes[8]; i++) {
+            conv1.add(
+                            Conv2d.builder()
+                                    .setKernelShape(new Shape(1, 1))
+                                    .setFilters(builder.filters[8])
+                                    .optStride(new Shape(builder.strides[8], builder.strides[8]))
+                                    .optBias(false)
+                                    .build())
+                    .add(
+                            BatchNorm.builder()
+                                    .optEpsilon(2E-5f)
+                                    .optMomentum(builder.batchNormMomentum)
+                                    .build())
+                    .add(Activation.reluBlock());
+        }
+
+        Block conv2 =
+                Conv2d.builder()
+                        .setKernelShape(new Shape(1, 1))
+                        .setFilters((int) builder.outSize)
+                        .optBias(false)
+                        .build();
+
+        return mobileNet
+                .add(pre)
+                .addAll(bottleNecks)
+                .add(conv1)
+                .add(Pool.globalAvgPool2dBlock())
+                .addSingleton(
+                        array -> array.reshape(array.getShape().get(0), builder.filters[8], 1, 1))
+                // reshape for conv1*1
+                .add(conv2)
+                // reshape for output
+                .addSingleton(array -> array.reshape(array.getShape().get(0), builder.outSize));
+    }
+
+    /**
+     * Creates a builder to build a {@link MobileNetV2}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** The Builder to construct a {@link MobileNetV2} object. */
+    public static final class Builder {
+
+        float batchNormMomentum = 0.9f;
+        long outSize = 10; // 10 as default for basic datasets like cifar-10 or mnist
+        int[] repeatTimes = {
+            1, 1, 2, 3, 4, 3, 3, 1, 1
+        }; // repeatTimes(n) of each block defined in the paper
+        int[] filters = {
+            32, 16, 24, 32, 64, 96, 160, 320, 1280
+        }; // filters(c) of each Block defined in the paper
+        int[] strides = {
+            2, 1, 2, 2, 2, 1, 2, 1, 1
+        }; // strides(s) of each block defined in the paper
+        int[] multiTimes = {
+            1, 6, 6, 6, 6, 6, 6
+        }; // multipleTimes(t) of each linearBottleneck defined in the paper
+
+        Builder() {}
+
+        /**
+         * 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;
+        }
+
+        /**
+         * Sets the filters(the value c defined in the paper) of customized MobileNetV2.
+         *
+         * @param filters the customized filter
+         * @return this {@code Builder}
+         */
+        public Builder optFilters(int[] filters) {
+            if (filters.length != FILTERLENGTH) {
+                throw new IllegalArgumentException(
+                        String.format(
+                                "optFilters requires filters of length %d, but was given filters of"
+                                        + " length %d instead",
+                                FILTERLENGTH, filters.length));
+            }
+            this.filters = filters;
+            return this;
+        }
+
+        /**
+         * Sets the repeatTimes(the value n defined in the paper) of each block of MobileNetV2.
+         *
+         * @param repeatTimes the customized repeatTimes
+         * @return this {@code Builder}
+         */
+        public Builder optRepeatTimes(int[] repeatTimes) {
+            if (repeatTimes.length != REPEATLENGTH) {
+                throw new IllegalArgumentException(
+                        String.format(
+                                "optRepeatTimes requires repeatTimes of length %d, but was given"
+                                        + " repeatTimes of length %d instead",
+                                REPEATLENGTH, repeatTimes.length));
+            }
+            this.repeatTimes = repeatTimes;
+            return this;
+        }
+
+        /**
+         * Sets the strides(the value s defined in the paper) of each block of MobileNetV2.
+         *
+         * @param strides the customized strides
+         * @return this {@code Builder}
+         */
+        public Builder optStrides(int[] strides) {
+            if (strides.length != STRIDELENGTH) {
+                throw new IllegalArgumentException(
+                        String.format(
+                                "optStrides requires strides of length %d, but was given strides of"
+                                        + " length %d instead",
+                                STRIDELENGTH, strides.length));
+            }
+            this.strides = strides;
+            return this;
+        }
+
+        /**
+         * Sets the multiTimes(the value t defined in the paper) of each bottleNeck of MobileNetV2.
+         *
+         * @param multiTimes the customized multiTimes
+         * @return this {@code Builder}
+         */
+        public Builder optMultiTimes(int[] multiTimes) {
+            if (multiTimes.length != MULTILENGTH) {
+                throw new IllegalArgumentException(
+                        String.format(
+                                "optMultiTimes requires multiTimes of length %d, but was given"
+                                        + " multiTimes of length %d instead",
+                                MULTILENGTH, multiTimes.length));
+            }
+            this.multiTimes = multiTimes;
+            return this;
+        }
+
+        /**
+         * Builds a {@link MobileNetV2} block.
+         *
+         * @return the {@link MobileNetV2} block
+         */
+        public Block build() {
+            return mobilenetV2(this);
+        }
+    }
+}


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/MobileNetV2Test.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/MobileNetV2Test.java
new file mode 100644
index 000000000..8fc9baa59
--- /dev/null
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/classification/MobileNetV2Test.java
@@ -0,0 +1,92 @@
+/*
+ * 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.MobileNetV2;
+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 MobileNetV2Test {
+    @Test
+    public void testTrain() {
+        TrainingConfig config =
+                new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
+                        .optDevices(Engine.getInstance().getDevices(2))
+                        .optInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+        Block mobilenet = MobileNetV2.builder().setOutSize(10).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.randomUniform(0, 1, 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();
+
+                NDArray expectedAtIndex0 =
+                        manager.ones(new Shape(32, 3, 1, 1)); // 32*3*1*1 for first layer
+                NDArray expectedAtIndex5 =
+                        manager.ones(new Shape(32, 32, 1, 1)); // 32*32*1*1 for pointWiseLayer1
+                NDArray expectedAtIndex10 =
+                        manager.ones(new Shape(32, 1, 3, 3)); // 32*1*3*3 for depthWiseLayer1
+                NDArray expectedAtIndex260 =
+                        manager.ones(
+                                new Shape(
+                                        1280, 320, 1,
+                                        1)); // 1280*320*1*1  for pointWiseLayer at last
+
+                Assertions.assertAlmostEquals(
+                        parameters.get(0).getValue().getArray(), expectedAtIndex0);
+                Assertions.assertAlmostEquals(
+                        parameters.get(5).getValue().getArray(), expectedAtIndex5);
+                Assertions.assertAlmostEquals(
+                        parameters.get(10).getValue().getArray(), expectedAtIndex10);
+                Assertions.assertAlmostEquals(
+                        parameters.get(260).getValue().getArray(), expectedAtIndex260);
+            }
+        }
+    }
+}