Repository: deepjavalibrary/djl
Pull request #2054 (2022-09-30 23:16:14)

Fix ames and various tabular improvements
First for Ames, this fixes a bug (misclassified column type) to fix ames. It also simplifies the Ames dataset by removing the featureToMap and replacing it with the standard usage of the categorical Featurizer.

For tabular, it moves the validation function from the builder to the constructor for tabular. This fixes the problem of how easy it is to forget the validation when creating a dataset. As proof of this, several datasets were missing the calls to validate. It adds two new tests for all features and no features in Ames.


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2054?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Base: **72.08**% // Head: **69.73**% // Decreases project coverage by **`-2.35%`** :warning:
> Coverage data is based on head [(`23ff41b`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2054?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) compared to base [(`bb5073f`)](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).
> Patch coverage: 68.85% of modified lines in pull request are covered.

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff              @@
##             master    #2054      +/-   ##
============================================
- Coverage     72.08%   69.73%   -2.36%     
- Complexity     5126     5982     +856     
============================================
  Files           473      600     +127     
  Lines         21970    26533    +4563     
  Branches       2351     2882     +531     
============================================
+ Hits          15838    18502    +2664     
- Misses         4925     6647    +1722     
- Partials       1207     1384     +177     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2054?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/2054/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/2054/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/2054/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/2054/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/2054/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/2054/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/2054/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: |
| [.../cv/translator/InstanceSegmentationTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2054/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yLmphdmE=) | `0.00% <0.00%> (-86.59%)` | :arrow_down: |
| [...nslator/InstanceSegmentationTranslatorFactory.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2054/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW5zdGFuY2VTZWdtZW50YXRpb25UcmFuc2xhdG9yRmFjdG9yeS5qYXZh) | `7.14% <0.00%> (-11.04%)` | :arrow_down: |
| [.../cv/translator/SemanticSegmentationTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2054/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%> (ø)` | |
| ... and [513 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2054/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) | |

Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary)

</details>

[:umbrella: View full report at Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2054?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).


Patch:
diff --git a/android/build.gradle b/android/build.gradle
index cc13cf3c0..b2a502542 100644
--- a/android/build.gradle
+++ b/android/build.gradle
@@ -6,7 +6,7 @@ buildscript {
     }
 
     dependencies {
-        classpath 'com.android.tools.build:gradle:4.2.2'
+        classpath 'com.android.tools.build:gradle:7.2.2'
     }
 }
 
diff --git a/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java b/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java
index 87e2f6485..7017370dd 100644
--- a/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java
+++ b/android/core/src/main/java/ai/djl/android/core/BitmapImageFactory.java
@@ -120,6 +120,13 @@ public class BitmapImageFactory extends ImageFactory {
         return new BitMapWrapper(bitmap);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public Image fromPixels(int[] pixels, int width, int height) {
+        Bitmap  bitmap = Bitmap.createBitmap(pixels, width, height, Bitmap.Config.ARGB_8888);
+        return new BitMapWrapper(bitmap);
+    }
+
     static class BitMapWrapper implements Image {
         private Bitmap bitmap;
 
@@ -141,10 +148,39 @@ public class BitmapImageFactory extends ImageFactory {
 
         /** {@inheritDoc} */
         @Override
-        public Object getWrappedImage() {
+        public Bitmap getWrappedImage() {
             return bitmap;
         }
 
+        /** {@inheritDoc} */
+        @Override
+        public BitMapWrapper resize(int width, int height, boolean copy) {
+            if (!copy && bitmap.getWidth() == width && bitmap.getHeight() == height) {
+                return this;
+            }
+            return new BitMapWrapper(Bitmap.createScaledBitmap(bitmap, width, height, true));
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Image getMask(int[][] mask) {
+            int w = mask[0].length;
+            int h = mask.length;
+            BitMapWrapper resized = resize(w, h, true);
+            Bitmap img = resized.getWrappedImage();
+            int[] pixels = new int[w * h];
+            int index = 0;
+            for (int y = 0; y < h; ++y) {
+                for (int x = 0; x < w; ++x) {
+                    if (mask[y][x] != 0) {
+                        pixels[index] = img.getPixel(x, y);
+                    }
+                    index++;
+                }
+            }
+            return new BitMapWrapper(Bitmap.createBitmap(pixels, w, h, Bitmap.Config.ARGB_8888));
+        }
+
         /** {@inheritDoc} */
         @Override
         public Image getSubImage(int x, int y, int w, int h) {
@@ -271,6 +307,24 @@ public class BitmapImageFactory extends ImageFactory {
             oldBitmap.recycle();
         }
 
+        /** {@inheritDoc} */
+        @Override
+        public void drawImage(Image overlay, boolean resize) {
+            if (!(overlay.getWrappedImage() instanceof Bitmap)) {
+                throw new IllegalArgumentException("Only Bitmap allowed");
+            }
+            if (resize) {
+                overlay = overlay.resize(getWidth(), getHeight(), false);
+            }
+            Bitmap target = Bitmap.createBitmap(getWidth(), getHeight(), Bitmap.Config.ARGB_8888);
+            Canvas canvas = new Canvas(target);
+            canvas.drawBitmap(bitmap, 0, 0, null);
+            canvas.drawBitmap((Bitmap) overlay.getWrappedImage(), 0, 0, null);
+            Bitmap oldBitmap = bitmap;
+            bitmap = target;
+            oldBitmap.recycle();
+        }
+
         private int randomColor() {
             return Color.rgb(
                     RandomUtils.nextInt(255), RandomUtils.nextInt(255), RandomUtils.nextInt(255));
diff --git a/android/gradle.properties b/android/gradle.properties
index 5179882b5..46e91c8cd 100644
--- a/android/gradle.properties
+++ b/android/gradle.properties
@@ -17,5 +17,5 @@ org.gradle.jvmargs=-Xmx1536m
 android.useAndroidX=true
 # Automatically convert third-party libraries to use AndroidX
 android.enableJetifier=true
-djl_version=0.19.0
+djl_version=0.20.0
 pytorch_version=1.12.1
diff --git a/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
index a66050657..d850b08a1 100644
--- a/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/BufferedImageFactory.java
@@ -127,6 +127,14 @@ public class BufferedImageFactory extends ImageFactory {
         return new BufferedImageWrapper(image);
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public Image fromPixels(int[] pixels, int width, int height) {
+        BufferedImage image = new BufferedImage(width, height, BufferedImage.TYPE_INT_ARGB);
+        image.setRGB(0, 0, width, height, pixels, 0, width);
+        return new BufferedImageWrapper(image);
+    }
+
     protected void save(BufferedImage image, OutputStream os, String type) throws IOException {
         ImageIO.write(image, type, os);
     }
@@ -153,10 +161,26 @@ public class BufferedImageFactory extends ImageFactory {
 
         /** {@inheritDoc} */
         @Override
-        public Object getWrappedImage() {
+        public BufferedImage getWrappedImage() {
             return image;
         }
 
+        /** {@inheritDoc} */
+        @Override
+        public BufferedImageWrapper resize(int width, int height, boolean copy) {
+            if (!copy && image.getWidth() == width && image.getHeight() == height) {
+                return this;
+            }
+
+            java.awt.Image img =
+                    image.getScaledInstance(width, height, java.awt.Image.SCALE_SMOOTH);
+            BufferedImage scaled = new BufferedImage(width, height, image.getType());
+            Graphics2D g2d = scaled.createGraphics();
+            g2d.drawImage(img, 0, 0, null);
+            g2d.dispose();
+            return new BufferedImageWrapper(scaled);
+        }
+
         /** {@inheritDoc} */
         @Override
         public Image getSubImage(int x, int y, int w, int h) {
@@ -174,6 +198,26 @@ public class BufferedImageFactory extends ImageFactory {
             return new BufferedImageWrapper(copy);
         }
 
+        /** {@inheritDoc} */
+        @Override
+        public Image getMask(int[][] mask) {
+            int w = mask[0].length;
+            int h = mask.length;
+            BufferedImageWrapper resized = resize(w, h, true);
+            BufferedImage img = resized.getWrappedImage();
+            int[] pixels = new int[w * h];
+            int index = 0;
+            for (int y = 0; y < h; ++y) {
+                for (int x = 0; x < w; ++x) {
+                    if (mask[y][x] != 0) {
+                        pixels[index] = img.getRGB(x, y);
+                    }
+                    index++;
+                }
+            }
+            return fromPixels(pixels, w, h);
+        }
+
         private void convertIdNeeded() {
             if (image.getType() == BufferedImage.TYPE_INT_ARGB) {
                 return;
@@ -308,6 +352,24 @@ public class BufferedImageFactory extends ImageFactory {
             g.dispose();
         }
 
+        /** {@inheritDoc} */
+        @Override
+        public void drawImage(Image overlay, boolean resize) {
+            if (!(overlay.getWrappedImage() instanceof BufferedImage)) {
+                throw new IllegalArgumentException("Only BufferedImage allowed");
+            }
+            if (resize) {
+                overlay = overlay.resize(getWidth(), getHeight(), false);
+            }
+            BufferedImage target =
+                    new BufferedImage(getWidth(), getHeight(), BufferedImage.TYPE_INT_ARGB);
+            Graphics2D g = (Graphics2D) target.getGraphics();
+            g.drawImage(image, 0, 0, null);
+            g.drawImage((BufferedImage) overlay.getWrappedImage(), 0, 0, null);
+            g.dispose();
+            image = target;
+        }
+
         private Color randomColor() {
             return new Color(RandomUtils.nextInt(255));
         }
diff --git a/api/src/main/java/ai/djl/modality/cv/Image.java b/api/src/main/java/ai/djl/modality/cv/Image.java
index 6e93aba69..26e7fe4ce 100644
--- a/api/src/main/java/ai/djl/modality/cv/Image.java
+++ b/api/src/main/java/ai/djl/modality/cv/Image.java
@@ -49,6 +49,24 @@ public interface Image {
      */
     Object getWrappedImage();
 
+    /**
+     * Creates a new resized image.
+     *
+     * @param width the new image's desired width
+     * @param height the new image's desired height
+     * @param copy false to return original image if size is the same
+     * @return the new resized image.
+     */
+    Image resize(int width, int height, boolean copy);
+
+    /**
+     * Returns a new {@code Image} of masked area.
+     *
+     * @param mask the mask for each pixel in the image
+     * @return the mask image.
+     */
+    Image getMask(int[][] mask);
+
     /**
      * Gets the subimage defined by a specified rectangular region.
      *
@@ -116,6 +134,14 @@ public interface Image {
      */
     void drawJoints(Joints joints);
 
+    /**
+     * Draws the overlay on the image.
+     *
+     * @param overlay the overlay image
+     * @param resize true to resize the overlay image to match the image
+     */
+    void drawImage(Image overlay, boolean resize);
+
     /** Flag indicates the color channel options for images. */
     enum Flag {
         GRAYSCALE,
diff --git a/api/src/main/java/ai/djl/modality/cv/ImageFactory.java b/api/src/main/java/ai/djl/modality/cv/ImageFactory.java
index 5e73d7ae6..17167a8b6 100644
--- a/api/src/main/java/ai/djl/modality/cv/ImageFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/ImageFactory.java
@@ -138,4 +138,14 @@ public abstract class ImageFactory {
      * @return {@link Image}
      */
     public abstract Image fromNDArray(NDArray array);
+
+    /**
+     * Gets {@link Image} from array.
+     *
+     * @param pixels the array of ARGB values used to initialize the pixels.
+     * @param width the width of the image
+     * @param height the height of the image
+     * @return {@link Image}
+     */
+    public abstract Image fromPixels(int[] pixels, int width, int height);
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/output/CategoryMask.java b/api/src/main/java/ai/djl/modality/cv/output/CategoryMask.java
new file mode 100644
index 000000000..4cdc75e21
--- /dev/null
+++ b/api/src/main/java/ai/djl/modality/cv/output/CategoryMask.java
@@ -0,0 +1,206 @@
+/*
+ * 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.modality.cv.output;
+
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.ImageFactory;
+import ai.djl.util.JsonSerializable;
+import ai.djl.util.JsonUtils;
+import ai.djl.util.RandomUtils;
+
+import com.google.gson.Gson;
+import com.google.gson.JsonElement;
+import com.google.gson.JsonSerializationContext;
+import com.google.gson.JsonSerializer;
+
+import java.lang.reflect.Type;
+import java.nio.ByteBuffer;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+
+/**
+ * A class representing the segmentation result of an image in an {@link
+ * ai.djl.Application.CV#SEMANTIC_SEGMENTATION} case.
+ */
+public class CategoryMask implements JsonSerializable {
+
+    private static final long serialVersionUID = 1L;
+
+    private static final int COLOR_BLACK = 0xFF000000;
+
+    private static final Gson GSON =
+            JsonUtils.builder()
+                    .registerTypeAdapter(CategoryMask.class, new SegmentationSerializer())
+                    .create();
+
+    private transient List<String> classes;
+    private int[][] mask;
+
+    /**
+     * Constructs a Mask with the given data.
+     *
+     * @param classes the list of classes
+     * @param mask the category mask for each pixel in the image
+     */
+    public CategoryMask(List<String> classes, int[][] mask) {
+        this.classes = classes;
+        this.mask = mask;
+    }
+
+    /**
+     * Returns the list of classes.
+     *
+     * @return list of classes
+     */
+    public List<String> getClasses() {
+        return classes;
+    }
+
+    /**
+     * Returns the class for each pixel.
+     *
+     * @return the class for each pixel
+     */
+    public int[][] getMask() {
+        return mask;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public ByteBuffer toByteBuffer() {
+        return ByteBuffer.wrap(toJson().getBytes(StandardCharsets.UTF_8));
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public String toJson() {
+        return GSON.toJson(this) + '\n';
+    }
+
+    /**
+     * Extracts the detected objects from the image.
+     *
+     * @param image the original image
+     * @return the detected objects from the image
+     */
+    public Image getMaskImage(Image image) {
+        return image.getMask(mask);
+    }
+
+    /**
+     * Extracts the specified object from the image.
+     *
+     * @param image the original image
+     * @param classId the class to extract from the image
+     * @return the specific object on the image
+     */
+    public Image getMaskImage(Image image, int classId) {
+        int width = mask[0].length;
+        int height = mask.length;
+        int[][] selected = new int[height][width];
+        for (int h = 0; h < height; ++h) {
+            for (int w = 0; w < width; ++w) {
+                selected[h][w] = mask[h][w] == classId ? 1 : 0;
+            }
+        }
+        return image.getMask(selected);
+    }
+
+    /**
+     * Extracts the background from the image.
+     *
+     * @param image the original image
+     * @return the background of the image
+     */
+    public Image getBackgroundImage(Image image) {
+        return getMaskImage(image, 0);
+    }
+
+    /**
+     * Highlights the detected object on the image with random colors.
+     *
+     * @param image the original image
+     * @param transparency the transparency of the overlay. Value is between 0 and 255 inclusive,
+     *     where 0 means the overlay is completely transparent and 255 means the overlay is
+     *     completely opaque.
+     */
+    public void drawMask(Image image, int transparency) {
+        drawMask(image, transparency, COLOR_BLACK);
+    }
+
+    /**
+     * Highlights the detected object on the image with random colors.
+     *
+     * @param image the original image
+     * @param transparency the transparency of the overlay. Value is between 0 and 255 inclusive,
+     *     where 0 means the overlay is completely opaque and 255 means the overlay is completely
+     *     transparent.
+     * @param background replace the background with specified background color, use transparent
+     *     color to remove background
+     */
+    public void drawMask(Image image, int transparency, int background) {
+        int[] colors = generateColors(background, transparency);
+        Image maskImage = getColorOverlay(colors);
+        image.drawImage(maskImage, true);
+    }
+
+    /**
+     * Highlights the specified object with specific color.
+     *
+     * @param image the original image
+     * @param classId the class to draw on the image
+     * @param color the rgb color with transparency
+     * @param transparency the transparency of the overlay
+     */
+    public void drawMask(Image image, int classId, int color, int transparency) {
+        int[] colors = new int[classes.size()];
+        colors[classId] = color & 0xFFFFFF | transparency << 24;
+        Image colorOverlay = getColorOverlay(colors);
+        image.drawImage(colorOverlay, true);
+    }
+
+    private Image getColorOverlay(int[] colors) {
+        int height = mask.length;
+        int width = mask[0].length;
+        int[] pixels = new int[width * height];
+        for (int h = 0; h < height; h++) {
+            for (int w = 0; w < width; w++) {
+                int index = mask[h][w];
+                pixels[h * width + w] = colors[index];
+            }
+        }
+        return ImageFactory.getInstance().fromPixels(pixels, width, height);
+    }
+
+    private int[] generateColors(int background, int transparency) {
+        int[] colors = new int[classes.size()];
+        colors[0] = background;
+        for (int i = 1; i < classes.size(); i++) {
+            int red = RandomUtils.nextInt(256);
+            int green = RandomUtils.nextInt(256);
+            int blue = RandomUtils.nextInt(256);
+            colors[i] = transparency << 24 | red << 16 | green << 8 | blue;
+        }
+        return colors;
+    }
+
+    /** A customized Gson serializer to serialize the {@code Segmentation} object. */
+    public static final class SegmentationSerializer implements JsonSerializer<CategoryMask> {
+
+        /** {@inheritDoc} */
+        @Override
+        public JsonElement serialize(CategoryMask src, Type type, JsonSerializationContext ctx) {
+            return ctx.serialize(src.getMask());
+        }
+    }
+}
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java
index b9e0b5726..9ca5942a7 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslator.java
@@ -13,33 +13,31 @@
 package ai.djl.modality.cv.translator;
 
 import ai.djl.modality.cv.Image;
-import ai.djl.modality.cv.ImageFactory;
-import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.output.CategoryMask;
 import ai.djl.modality.cv.util.NDImageUtils;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
-import ai.djl.ndarray.NDManager;
-import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.translate.ArgumentsUtil;
 import ai.djl.translate.Transform;
+import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorContext;
-import ai.djl.util.RandomUtils;
 
-import java.nio.ByteBuffer;
+import java.io.IOException;
+import java.util.List;
 import java.util.Map;
 
 /**
- * A {@link BaseImageTranslator} that post-process the {@link NDArray} into {@link DetectedObjects}
- * with boundaries at the detailed pixel level.
+ * A {@link Translator} that post-process the {@link Image} into {@link CategoryMask} with output
+ * mask representing the class that each pixel in the original image belong to.
  */
-public class SemanticSegmentationTranslator extends BaseImageTranslator<Image> {
+public class SemanticSegmentationTranslator extends BaseImageTranslator<CategoryMask> {
 
+    private SynsetLoader synsetLoader;
     private final int shortEdge;
     private final int maxEdge;
 
-    private static final int CHANNEL = 3;
-    private static final int CLASSNUM = 21;
+    private List<String> classes;
 
     /**
      * Creates the Semantic Segmentation translator from the given builder.
@@ -48,78 +46,58 @@ public class SemanticSegmentationTranslator extends BaseImageTranslator<Image> {
      */
     public SemanticSegmentationTranslator(Builder builder) {
         super(builder);
+        this.synsetLoader = builder.synsetLoader;
         this.shortEdge = builder.shortEdge;
         this.maxEdge = builder.maxEdge;
 
         pipeline.insert(0, null, new ResizeShort());
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(TranslatorContext ctx) throws IOException {
+        if (classes == null) {
+            classes = synsetLoader.load(ctx.getModel());
+        }
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDList processInput(TranslatorContext ctx, Image image) {
-        ctx.setAttachment("originalHeight", image.getHeight());
-        ctx.setAttachment("originalWidth", image.getWidth());
         return super.processInput(ctx, image);
     }
 
     /** {@inheritDoc} */
     @Override
-    public Image processOutput(TranslatorContext ctx, NDList list) {
+    public CategoryMask processOutput(TranslatorContext ctx, NDList list) {
         // scores contains the probabilities of each pixel being a certain object
         float[] scores = list.get(1).toFloatArray();
         Shape shape = list.get(1).getShape();
         int width = (int) shape.get(2);
         int height = (int) shape.get(1);
-
-        // build image array
-        try (NDManager manager = NDManager.newBaseManager()) {
-            int imageSize = width * height;
-            ByteBuffer bb = manager.allocateDirect(CHANNEL * imageSize);
-            int r = 0; // adjustment for red pixel
-            int g = 1; // adjustment for green pixel
-            int b = 2; // adjustment for blue pixel
-            byte[][] colors = new byte[CLASSNUM][3];
-            for (int i = 0; i < CLASSNUM; i++) {
-                byte red = (byte) RandomUtils.nextInt(256);
-                byte green = (byte) RandomUtils.nextInt(256);
-                byte blue = (byte) RandomUtils.nextInt(256);
-                colors[i][r] = red;
-                colors[i][g] = green;
-                colors[i][b] = blue;
-            }
-
-            // change color of pixels in image array where objects have been detected
-            for (int h = 0; h < height; h++) {
-                for (int w = 0; w < width; w++) {
-                    int index = h * width + w;
-                    int maxi = 0;
-                    double maxnum = -Double.MAX_VALUE;
-                    for (int i = 0; i < CLASSNUM; i++) {
-                        // get score for each i at the h,w pixel of the image
-                        float score = scores[i * (imageSize) + index];
-                        if (score > maxnum) {
-                            maxnum = score;
-                            maxi = i;
-                        }
-                    }
-                    if (maxi > 0) {
-                        bb.put(colors[maxi][r]);
-                        bb.put(colors[maxi][g]);
-                        bb.put(colors[maxi][b]);
-                    } else {
-                        bb.position(bb.position() + 3);
+        int[][] mask = new int[height][width];
+
+        int imageSize = width * height;
+
+        // Build mask array
+        int numOfClasses = classes.size();
+        for (int h = 0; h < height; h++) {
+            for (int w = 0; w < width; w++) {
+                int index = h * width + w;
+                int maxi = 0;
+                double maxnum = -Double.MAX_VALUE;
+                for (int i = 0; i < numOfClasses; ++i) {
+                    // get score for each i at the h,w pixel of the image
+                    float score = scores[i * imageSize + index];
+                    if (score > maxnum) {
+                        maxnum = score;
+                        maxi = i;
                     }
                 }
+                mask[h][w] = maxi;
             }
-            bb.rewind();
-            int originW = (int) ctx.getAttachment("originalWidth");
-            int originH = (int) ctx.getAttachment("originalHeight");
-            NDArray fullImage =
-                    manager.create(bb, new Shape(height, width, CHANNEL), DataType.UINT8);
-            NDArray resized = NDImageUtils.resize(fullImage, originW, originH);
-
-            return ImageFactory.getInstance().fromNDArray(resized);
         }
+        return new CategoryMask(classes, mask);
     }
 
     /**
diff --git a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java
index ec69dabdd..6374fc1fa 100644
--- a/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java
+++ b/api/src/main/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactory.java
@@ -13,20 +13,13 @@
 package ai.djl.modality.cv.translator;
 
 import ai.djl.Model;
-import ai.djl.modality.Input;
-import ai.djl.modality.Output;
 import ai.djl.modality.cv.Image;
-import ai.djl.modality.cv.translator.wrapper.FileTranslator;
-import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
-import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.modality.cv.output.CategoryMask;
 import ai.djl.translate.Translator;
 import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.Pair;
 
-import java.io.InputStream;
 import java.lang.reflect.Type;
-import java.net.URL;
-import java.nio.file.Path;
 import java.util.HashSet;
 import java.util.Map;
 import java.util.Set;
@@ -37,7 +30,7 @@ public class SemanticSegmentationTranslatorFactory implements TranslatorFactory
     private static final Set<Pair<Type, Type>> SUPPORTED_TYPES = new HashSet<>();
 
     static {
-        SUPPORTED_TYPES.add(new Pair<>(Image.class, Image.class));
+        SUPPORTED_TYPES.add(new Pair<>(Image.class, CategoryMask.class));
     }
 
     /** {@inheritDoc} */
@@ -45,18 +38,8 @@ public class SemanticSegmentationTranslatorFactory implements TranslatorFactory
     @SuppressWarnings("unchecked")
     public <I, O> Translator<I, O> newInstance(
             Class<I> input, Class<O> output, Model model, Map<String, ?> arguments) {
-        SemanticSegmentationTranslator translator =
-                SemanticSegmentationTranslator.builder(arguments).build();
-        if (input == Image.class && output == Image.class) {
-            return (Translator<I, O>) translator;
-        } else if (input == Path.class && output == Image.class) {
-            return (Translator<I, O>) new FileTranslator<>(translator);
-        } else if (input == URL.class && output == Image.class) {
-            return (Translator<I, O>) new UrlTranslator<>(translator);
-        } else if (input == InputStream.class && output == Image.class) {
-            return (Translator<I, O>) new InputStreamTranslator<>(translator);
-        } else if (input == Input.class && output == Output.class) {
-            return (Translator<I, O>) new ImageServingTranslator(translator);
+        if (input == Image.class && output == CategoryMask.class) {
+            return (Translator<I, O>) SemanticSegmentationTranslator.builder(arguments).build();
         }
         throw new IllegalArgumentException("Unsupported input/output types.");
     }
diff --git a/api/src/main/java/ai/djl/training/loss/TabNetLoss.java b/api/src/main/java/ai/djl/training/loss/TabNetClassificationLoss.java
similarity index 53%
rename from api/src/main/java/ai/djl/training/loss/TabNetLoss.java
rename to api/src/main/java/ai/djl/training/loss/TabNetClassificationLoss.java
index 296f4b93d..58756054d 100644
--- a/api/src/main/java/ai/djl/training/loss/TabNetLoss.java
+++ b/api/src/main/java/ai/djl/training/loss/TabNetClassificationLoss.java
@@ -16,29 +16,32 @@ import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 
 /**
- * Calculates the loss for tabNet. Actually, loss has been calculated through the forward function
- * of tabNet. What's done here is just getting the loss function from prediction.
+ * Calculates the loss for tabNet in Classification tasks.
+ *
+ * <p>Actually, tabNet is not only used for Supervised Learning, it's also widely used in
+ * unsupervised learning. For unsupervised learning, it should come from the decoder(aka
+ * attentionTransformer of tabNet)
  */
-public class TabNetLoss extends Loss {
-    /** Calculates the loss of a TabNet instance. */
-    public TabNetLoss() {
-        this("TabNetLoss");
+public final class TabNetClassificationLoss extends Loss {
+    /** Calculates the loss of a TabNet instance for regression tasks. */
+    public TabNetClassificationLoss() {
+        this("TabNetClassificationLoss");
     }
 
     /**
-     * Calculates the loss of a TabNet instance.
+     * Calculates the loss of a TabNet instance for regression tasks.
      *
      * @param name the name of the loss function
      */
-    public TabNetLoss(String name) {
+    public TabNetClassificationLoss(String name) {
         super(name);
     }
 
     /** {@inheritDoc} */
     @Override
     public NDArray evaluate(NDList labels, NDList predictions) {
-        // loss is already calculated inside the forward of tabNet
-        // so here we just need to get it out from prediction
-        return predictions.get(1);
+        return Loss.softmaxCrossEntropyLoss()
+                .evaluate(labels, new NDList(predictions.get(0)))
+                .add(predictions.get(1));
     }
 }
diff --git a/api/src/main/java/ai/djl/training/loss/TabNetRegressionLoss.java b/api/src/main/java/ai/djl/training/loss/TabNetRegressionLoss.java
new file mode 100644
index 000000000..21f290136
--- /dev/null
+++ b/api/src/main/java/ai/djl/training/loss/TabNetRegressionLoss.java
@@ -0,0 +1,51 @@
+/*
+ * 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.training.loss;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+
+/**
+ * Calculates the loss of tabNet for regression tasks.
+ *
+ * <p>Actually, tabNet is not only used for Supervised Learning, it's also widely used in
+ * unsupervised learning. For unsupervised learning, it should come from the decoder(aka
+ * attentionTransformer of tabNet)
+ */
+public class TabNetRegressionLoss extends Loss {
+    /** Calculates the loss of a TabNet instance for regression tasks. */
+    public TabNetRegressionLoss() {
+        this("TabNetRegressionLoss");
+    }
+
+    /**
+     * Calculates the loss of a TabNet instance for regression tasks.
+     *
+     * @param name the name of the loss function
+     */
+    public TabNetRegressionLoss(String name) {
+        super(name);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray evaluate(NDList labels, NDList predictions) {
+        // sparseLoss is already calculated inside the forward of tabNet
+        // so here we just need to get it out from prediction
+        return labels.singletonOrThrow()
+                .sub(predictions.get(0))
+                .square()
+                .mean()
+                .add(predictions.get(1));
+    }
+}
diff --git a/api/src/main/java/ai/djl/training/loss/YOLOv3Loss.java b/api/src/main/java/ai/djl/training/loss/YOLOv3Loss.java
new file mode 100644
index 000000000..462b86e28
--- /dev/null
+++ b/api/src/main/java/ai/djl/training/loss/YOLOv3Loss.java
@@ -0,0 +1,487 @@
+/*
+ * 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.training.loss;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDArrays;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Activation;
+
+/**
+ * {@code YOLOv3Loss} is an implementation of {@link Loss}. It is used to compute the loss while
+ * training a YOLOv3 model for object detection. It involves computing the targets given the
+ * generated anchors, labels and predictions, and then computing the sum of class predictions and
+ * bounding box predictions.
+ */
+public final class YOLOv3Loss extends Loss {
+    // TODO: currently not finished, still have some bugs inside and it can only be trained with
+    // PyTorch Engine
+    /*
+       PRESETANCHORS shapes come from the K-means clustering of COCO dataset, which image size is 416*416
+       it can be reshaped into any shape like 256*256, just multiply each value with 256/416
+    */
+    private static final float[] PRESETANCHORS = {
+        116, 90, 156, 198, 373, 326,
+        30, 61, 62, 45, 59, 119,
+        10, 13, 16, 30, 33, 23
+    };
+
+    private float[] anchors;
+    private int numClasses;
+    private int boxAttr;
+    private Shape inputShape;
+    private float ignoreThreshold;
+    private NDManager manager;
+    private static final float EPSILON = 1e-7f;
+
+    private YOLOv3Loss(Builder builder) {
+        super(builder.name);
+        this.anchors = builder.anchorsArray;
+        this.numClasses = builder.numClasses;
+        this.boxAttr = builder.numClasses + 5; // 5 for x,y,h,w,c
+        this.inputShape = builder.inputShape;
+        this.ignoreThreshold = builder.ignoreThreshold;
+    }
+
+    /**
+     * Gets the preset anchors of YoloV3.
+     *
+     * @return the preset anchors of YoloV3
+     */
+    public static float[] getPresetAnchors() {
+        return PRESETANCHORS.clone();
+    }
+
+    /**
+     * Make the value of given NDArray between tMin and tMax.
+     *
+     * @param tList the given NDArray
+     * @param tMin the min value
+     * @param tMax the max value
+     * @return a NDArray where values are set between tMin and tMax
+     */
+    public NDArray clipByTensor(NDArray tList, float tMin, float tMax) {
+        NDArray result = tList.gte(tMin).mul(tList).add(tList.lt(tMin).mul(tMin));
+        result = result.lte(tMax).mul(result).add(result.gt(tMax).mul(tMax));
+        return result;
+    }
+
+    /**
+     * Calculates the MSELoss between prediction and target.
+     *
+     * @param prediction the prediction array
+     * @param target the target array
+     * @return the MSELoss between prediction and target
+     */
+    public NDArray mseLoss(NDArray prediction, NDArray target) {
+        return prediction.sub(target).pow(2);
+    }
+
+    /**
+     * Calculates the BCELoss between prediction and target.
+     *
+     * @param prediction the prediction array
+     * @param target the target array
+     * @return the BCELoss between prediction and target
+     */
+    public NDArray bceLoss(NDArray prediction, NDArray target) {
+        prediction = clipByTensor(prediction, EPSILON, (float) (1.0 - EPSILON));
+        return prediction
+                .log()
+                .mul(target)
+                .add(prediction.mul(-1).add(1).log().mul(target.mul(-1).add(1)))
+                .mul(-1);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray evaluate(NDList labels, NDList predictions) {
+        manager = predictions.getManager();
+
+        /*
+        three outputs in total
+        NDArray out0 = predictions.get(0),  //Shape = (batchSize * 75 * 13 * 13) 75 = 3*(20+5)
+                out1 = predictions.get(1),  //Shape = (batchSize * 75 * 26 * 26)
+                out2 = predictions.get(2);  //Shape = (batchSize * 75 * 52 * 52)
+        */
+
+        NDArray[] lossComponents = new NDArray[3];
+        for (int i = 0; i < 3; i++) {
+            lossComponents[i] = evaluateOneOutput(i, predictions.get(i), labels.singletonOrThrow());
+        }
+
+        // calculate the final loss
+        return NDArrays.add(lossComponents);
+    }
+
+    /**
+     * Computes the Loss for one outputLayer.
+     *
+     * @param componentIndex which outputLayer does current input represent. the shape should be
+     *     (13*13,26*26,52*52)
+     * @param input one prediction layer of YOLOv3
+     * @param labels target labels. Must contain (offsetLabels, masks, classlabels)
+     * @return the total loss of a outputLayer
+     */
+    public NDArray evaluateOneOutput(int componentIndex, NDArray input, NDArray labels) {
+        int batchSize = (int) input.getShape().get(0);
+        int inW = (int) input.getShape().get(2);
+        int inH = (int) input.getShape().get(3);
+
+        NDArray prediction =
+                input.reshape(batchSize, 3, boxAttr, inW, inH)
+                        .transpose(1, 0, 3, 4, 2); // reshape into (3,batchSize,inW,inH,attrs)
+
+        // the prediction value of x,y,w,h which shape should be (3,batchSize,inW,inH)
+        NDArray x = Activation.sigmoid(prediction.get("...,0"));
+        NDArray y = Activation.sigmoid(prediction.get("...,1"));
+        NDArray w = prediction.get("...,2");
+        NDArray h = prediction.get("...,3");
+
+        // Confidence of whether there is an object and conditional probability of each class
+        // it should be reshaped into (batchSize,3)
+        NDArray conf = Activation.sigmoid(prediction.get("...,4")).transpose(1, 0, 2, 3);
+        NDArray predClass = Activation.sigmoid(prediction.get("...,5:")).transpose(1, 0, 2, 3, 4);
+
+        // get an NDList of groundTruth which contains boxLossScale and groundTruth
+        NDList truthList = getTarget(labels, inH, inW);
+
+        /*
+           boxLossScale shape should be: (batchSize,3,inW,inH)
+           groundTruth shape should be: (3,batchSize,inW,inH,boxAttr)
+        */
+        NDArray boxLossScale = truthList.get(0).transpose(1, 0, 2, 3);
+        NDArray groundTruth = truthList.get(1);
+
+        // iou shape should be: (batchSize,3 ,inW,inH)
+        NDArray iou =
+                calculateIOU(x, y, groundTruth.get("...,0:4"), componentIndex)
+                        .transpose(1, 0, 2, 3);
+
+        // get noObjMask and objMask
+        NDArray noObjMask =
+                NDArrays.where(
+                        iou.lte(ignoreThreshold), manager.ones(iou.getShape()), manager.create(0f));
+
+        NDArray objMask = iou.argMax(1).oneHot(3).transpose(0, 3, 1, 2);
+
+        objMask =
+                NDArrays.where(
+                        iou.gte(ignoreThreshold / 2),
+                        objMask,
+                        manager.zeros(objMask.getShape())); // to get rid of wrong ones
+        noObjMask = NDArrays.where(objMask.eq(1f), manager.zeros(noObjMask.getShape()), noObjMask);
+
+        NDArray xTrue = groundTruth.get("...,0");
+        NDArray yTrue = groundTruth.get("...,1");
+        NDArray wTrue = groundTruth.get("...,2");
+        NDArray hTrue = groundTruth.get("...,3");
+        NDArray classTrue = groundTruth.get("...,4:").transpose(1, 0, 2, 3, 4);
+
+        NDArray widths =
+                manager.create(
+                                new float[] {
+                                    anchors[componentIndex * 6],
+                                    anchors[componentIndex * 6 + 2],
+                                    anchors[componentIndex * 6 + 4]
+                                })
+                        .div(inputShape.get(0));
+
+        NDArray heights =
+                manager.create(
+                                new float[] {
+                                    anchors[componentIndex * 6 + 1],
+                                    anchors[componentIndex * 6 + 3],
+                                    anchors[componentIndex * 6 + 5]
+                                })
+                        .div(inputShape.get(1));
+
+        // three loss parts: box Loss, confidence Loss, and class Loss
+        NDArray boxLoss =
+                objMask.mul(boxLossScale)
+                        .mul(
+                                NDArrays.add(
+                                                xTrue.sub(x).pow(2),
+                                                yTrue.sub(y).pow(2),
+                                                wTrue.sub(
+                                                                w.exp()
+                                                                        .mul(
+                                                                                widths.broadcast(
+                                                                                                inH,
+                                                                                                inW,
+                                                                                                batchSize,
+                                                                                                3)
+                                                                                        .transpose(
+                                                                                                3,
+                                                                                                2,
+                                                                                                1,
+                                                                                                0)))
+                                                        .pow(2),
+                                                hTrue.sub(
+                                                                h.exp()
+                                                                        .mul(
+                                                                                heights.broadcast(
+                                                                                                inH,
+                                                                                                inW,
+                                                                                                batchSize,
+                                                                                                3)
+                                                                                        .transpose(
+                                                                                                3,
+                                                                                                2,
+                                                                                                1,
+                                                                                                0)))
+                                                        .pow(2))
+                                        .transpose(1, 0, 2, 3))
+                        .sum();
+
+        NDArray confLoss =
+                objMask.mul(
+                                conf.add(EPSILON)
+                                        .log()
+                                        .mul(-1)
+                                        .add(bceLoss(predClass, classTrue).sum(new int[] {4})))
+                        .sum();
+
+        NDArray noObjLoss = noObjMask.mul(conf.mul(-1).add(1 + EPSILON).log().mul(-1)).sum();
+
+        return boxLoss.add(confLoss).add(noObjLoss).div(batchSize);
+    }
+
+    /**
+     * Gets target NDArray for a given evaluator.
+     *
+     * @param labels the true labels
+     * @param inH the height of current layer
+     * @param inW the width of current layer
+     * @return an NDList of {boxLossScale and groundTruth}
+     */
+    public NDList getTarget(NDArray labels, int inH, int inW) {
+        int batchSize = (int) labels.size(0);
+
+        // the loss Scale of a box, used to punctuate small boxes
+        NDList boxLossComponents = new NDList();
+
+        // the groundTruth of a true object in pictures
+        NDList groundTruthComponents = new NDList();
+
+        // Shape of labels:(batchSize,objectNum,5)
+
+        for (int batch = 0; batch < batchSize; batch++) {
+            if (labels.get(batch).size(0) == 0) {
+                continue; // no object in current picture
+            }
+
+            NDArray boxLoss = manager.zeros(new Shape(inW, inH), DataType.FLOAT32);
+            NDArray groundTruth = manager.zeros(new Shape(inW, inH, boxAttr - 1), DataType.FLOAT32);
+
+            NDArray picture = labels.get(batch);
+            // the shape should be (objectNums,5)
+            NDArray xgt = picture.get("...,1").add(picture.get("...,3").div(2)).mul(inW);
+            // Center of x should be X value in labels and add half of the width and multiplies the
+            // input width to get which grid cell it's in
+            NDArray ygt = picture.get("...,2").add(picture.get("...,4").div(2)).mul(inH);
+            // Center of y is the same as well
+            NDArray wgt = picture.get("...,3");
+            // the width of the ground truth box
+            NDArray hgt = picture.get("...,4");
+            // the height of the ground truth box
+            // we should transform the presentation of true class, like
+            // [[0],[1],[2]]->[[1,0,0,...0],[0,1,0,...,0],[0,0,1,...,0]]
+            NDArray objectClass = picture.get("...,0");
+            objectClass = objectClass.oneHot(numClasses);
+
+            NDArray curLabel = labels.get(batch); // curLabel shape:(objectNum,5)
+            int objectNum = (int) curLabel.size(0);
+
+            for (int i = 0; i < objectNum; i++) {
+                // for each object, the middle of the object(x and y) should be in one grid cell of
+                // 13*13
+                // the tx and ty should indicate the grid cell and bx and by should indicate the
+                // movement from top-left of the grid cell
+                int tx = (int) xgt.get(i).getFloat();
+                int ty = (int) ygt.get(i).getFloat();
+                float bx = xgt.get(i).getFloat() - tx;
+                float by = ygt.get(i).getFloat() - ty;
+
+                String index = tx + "," + ty;
+                // set groundTruth
+                groundTruth.set(new NDIndex(index + ",0"), bx);
+                groundTruth.set(new NDIndex(index + ",1"), by);
+                groundTruth.set(new NDIndex(index + ",2"), wgt.getFloat(i));
+                groundTruth.set(new NDIndex(index + ",3"), hgt.getFloat(i));
+                groundTruth.set(new NDIndex(index + ",4:"), objectClass.get(i));
+
+                // set boxLoss
+                boxLoss.set(new NDIndex(index), 2 - wgt.getFloat(i) * hgt.getFloat(i));
+            }
+            boxLossComponents.add(boxLoss);
+            groundTruthComponents.add(groundTruth);
+        }
+
+        NDArray boxLossScale = NDArrays.stack(boxLossComponents).broadcast(3, batchSize, inW, inH);
+        NDArray groundTruth =
+                NDArrays.stack(groundTruthComponents)
+                        .broadcast(3, batchSize, inW, inH, boxAttr - 1);
+
+        return new NDList(boxLossScale, groundTruth);
+    }
+
+    /**
+     * Calculates the IOU between priori Anchors and groundTruth.
+     *
+     * @param predx the tx value of prediction
+     * @param predy the ty value of prediction
+     * @param groundTruth the groundTruth value of labels
+     * @param componentIndex the current component Index
+     * @return an NDArray of IOU
+     */
+    public NDArray calculateIOU(
+            NDArray predx, NDArray predy, NDArray groundTruth, int componentIndex) {
+        int inW = (int) predx.getShape().get(2);
+        int inH = (int) predx.getShape().get(3);
+        int strideW = (int) inputShape.get(0) / inW;
+        int strideH = (int) inputShape.get(1) / inH;
+
+        NDList iouComponent = new NDList();
+        // shape of predx, predy should all be (3,batchSize,inW,inH)
+        // shape of groundTruth should be (3,batchSize,inW,inH,4)
+        for (int i = 0; i < 3; i++) {
+            NDArray curPredx = predx.get(i);
+            NDArray curPredy = predy.get(i);
+            float width = anchors[componentIndex * 6 + 2 * i] / strideW;
+            float height = anchors[componentIndex * 6 + 2 * i + 1] / strideH;
+            NDArray predLeft = curPredx.sub(width / 2);
+            NDArray predRight = curPredx.add(width / 2);
+            NDArray predTop = curPredy.sub(height / 2);
+            NDArray predBottom = curPredy.add(height / 2);
+
+            NDArray truth = groundTruth.get(i);
+
+            NDArray trueLeft = truth.get("...,0").sub(truth.get("...,2").mul(inW).div(2));
+            NDArray trueRight = truth.get("...,0").add(truth.get("...,2").mul(inW).div(2));
+            NDArray trueTop = truth.get("...,1").sub(truth.get("...,3").mul(inH).div(2));
+            NDArray trueBottom = truth.get("...,1").add(truth.get("...,3").mul(inH).div(2));
+
+            NDArray left = NDArrays.maximum(predLeft, trueLeft);
+            NDArray right = NDArrays.minimum(predRight, trueRight);
+            NDArray top = NDArrays.maximum(predTop, trueTop);
+            NDArray bottom = NDArrays.minimum(predBottom, trueBottom);
+
+            NDArray inter = right.sub(left).mul(bottom.sub(top));
+            NDArray union =
+                    truth.get("...,2")
+                            .mul(inW)
+                            .mul(truth.get("...,3").mul(inH))
+                            .add(width * height)
+                            .sub(inter)
+                            .add(EPSILON); // should not be divided by zero
+
+            iouComponent.add(inter.div(union));
+        }
+
+        return NDArrays.stack(iouComponent);
+    }
+
+    /**
+     * Creates a new builder to build a {@link YOLOv3Loss}.
+     *
+     * @return a new builder;
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** The Builder to construct a {@link YOLOv3Loss} object. */
+    public static class Builder {
+        private String name = "YOLOv3Loss";
+        private float[] anchorsArray = PRESETANCHORS;
+        private int numClasses = 20;
+        private Shape inputShape = new Shape(419, 419);
+        private float ignoreThreshold = 0.5f;
+
+        /**
+         * Sets the loss name of YoloV3Loss.
+         *
+         * @param name the name of loss function
+         * @return this {@code Builder}
+         */
+        public Builder setName(String name) {
+            this.name = name;
+            return this;
+        }
+
+        /**
+         * Sets the preset anchors for YoloV3.
+         *
+         * @param anchorsArray the anchors in float array
+         * @return this {@code Builder}
+         */
+        public Builder setAnchorsArray(float[] anchorsArray) {
+            if (anchorsArray.length != PRESETANCHORS.length) {
+                throw new IllegalArgumentException(
+                        String.format(
+                                "setAnchorsArray requires anchors of length %d, but was given"
+                                        + " filters of length %d instead",
+                                PRESETANCHORS.length, anchorsArray.length));
+            }
+            this.anchorsArray = anchorsArray;
+            return this;
+        }
+
+        /**
+         * Sets the number of total classes.
+         *
+         * @param numClasses the number of total classes
+         * @return this {@code Builder}
+         */
+        public Builder setNumClasses(int numClasses) {
+            this.numClasses = numClasses;
+            return this;
+        }
+
+        /**
+         * Sets the shape of the input picture.
+         *
+         * @param inputShape the shape of input picture.
+         * @return this {@code Builder}
+         */
+        public Builder setInputShape(Shape inputShape) {
+            this.inputShape = inputShape;
+            return this;
+        }
+
+        /**
+         * Sets the ignoreThreshold for iou to check if we think it detects a picture.
+         *
+         * @param ignoreThreshold the ignore threshold
+         * @return this {@code Builder}
+         */
+        public Builder optIgnoreThreshold(float ignoreThreshold) {
+            this.ignoreThreshold = ignoreThreshold;
+            return this;
+        }
+
+        /**
+         * Builds a {@link YOLOv3Loss} instance.
+         *
+         * @return a {@link YOLOv3Loss} instance.
+         */
+        public YOLOv3Loss build() {
+            return new YOLOv3Loss(this);
+        }
+    }
+}
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/BananaDetection.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/BananaDetection.java
index a9555d353..b86f92078 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/BananaDetection.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/BananaDetection.java
@@ -88,6 +88,12 @@ public class BananaDetection extends ObjectDetectionDataset {
         return new PairList<>(Collections.singletonList(labels.get((int) index)));
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public List<String> getClasses() {
+        return Collections.singletonList("banana");
+    }
+
     /** {@inheritDoc} */
     @Override
     protected long availableSize() {
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/CocoDetection.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/CocoDetection.java
index e45abd33e..1f36b4a5e 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/CocoDetection.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/CocoDetection.java
@@ -45,14 +45,13 @@ import java.util.Optional;
  * <p>Each image might have different {@link ai.djl.ndarray.types.Shape}s.
  */
 public class CocoDetection extends ObjectDetectionDataset {
-
+    // TODO: Add synset logic for coco dataset
     private static final String ARTIFACT_ID = "coco";
     private static final String VERSION = "1.0";
 
     private Usage usage;
     private List<Path> imagePaths;
     private List<PairList<Long, Rectangle>> labels;
-
     private MRL mrl;
     private boolean prepared;
 
@@ -79,6 +78,13 @@ public class CocoDetection extends ObjectDetectionDataset {
         return labels.get(Math.toIntExact(index));
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public List<String> getClasses() {
+        throw new UnsupportedOperationException(
+                "getClasses() for CocoDetection has not been implemented yet.");
+    }
+
     /** {@inheritDoc} */
     @Override
     public void prepare(Progress progress) throws IOException {
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/ObjectDetectionDataset.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/ObjectDetectionDataset.java
index c48dd669f..84c7eb423 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/ObjectDetectionDataset.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/ObjectDetectionDataset.java
@@ -20,13 +20,13 @@ import ai.djl.util.Pair;
 import ai.djl.util.PairList;
 
 import java.io.IOException;
+import java.util.List;
 
 /**
  * A helper to create {@link ai.djl.training.dataset.Dataset}s for {@link
  * ai.djl.Application.CV#OBJECT_DETECTION}.
  */
 public abstract class ObjectDetectionDataset extends ImageDataset {
-
     /**
      * Creates a new instance of {@link ObjectDetectionDataset} with the given necessary
      * configurations.
@@ -68,4 +68,11 @@ public abstract class ObjectDetectionDataset extends ImageDataset {
      * @throws IOException if the data could not be loaded
      */
     public abstract PairList<Long, Rectangle> getObjects(long index) throws IOException;
+
+    /**
+     * Returns the classes that detected objects in the dataset can be classified into.
+     *
+     * @return the classes that detected objects in the dataset can be classified into.
+     */
+    public abstract List<String> getClasses();
 }
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/PikachuDetection.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/PikachuDetection.java
index 770858c75..dfdcb175c 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/PikachuDetection.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/PikachuDetection.java
@@ -127,6 +127,12 @@ public class PikachuDetection extends ObjectDetectionDataset {
         return new PairList<>(Collections.singletonList(labels.get((int) index)));
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public List<String> getClasses() {
+        return Collections.singletonList("pikachu");
+    }
+
     /** {@inheritDoc} */
     @Override
     protected long availableSize() {
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/AmesRandomAccess.java b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/AmesRandomAccess.java
index b9b4e8edb..e0e8ee7cf 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/AmesRandomAccess.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/AmesRandomAccess.java
@@ -29,7 +29,6 @@ import java.io.Reader;
 import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
 import java.util.List;
-import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
 
@@ -44,6 +43,9 @@ import java.util.Set;
  * <p>Test Set: 1459 Records
  *
  * <p>Can enable/disable features Set one hot vector for categorical variables
+ *
+ * <p>Call {@link Builder#addAllFeatures()} to include all features from the dataset. The label is a
+ * numeric column named "saleprice".
  */
 public class AmesRandomAccess extends CsvDataset {
 
@@ -200,15 +202,29 @@ public class AmesRandomAccess extends CsvDataset {
         public Builder addFeature(String name, boolean onehotEncode) {
             parseFeatures();
             if (af.categorical.contains(name)) {
-                Map<String, Integer> map = af.featureToMap.get(name);
-                if (map == null) {
-                    return addCategoricalFeature(name, onehotEncode);
-                }
-                return addCategoricalFeature(name, map, onehotEncode);
+                return addCategoricalFeature(name, onehotEncode);
             }
             return addNumericFeature(name);
         }
 
+        /**
+         * Adds all features to the features set.
+         *
+         * @return this builder
+         */
+        public Builder addAllFeatures() {
+            if (features.isEmpty()) {
+                parseFeatures();
+                for (String name : af.featureArray) {
+                    addFeature(name);
+                }
+            }
+            if (labels.isEmpty()) {
+                addNumericLabel("saleprice");
+            }
+            return this;
+        }
+
         /**
          * Returns the available features of this dataset.
          *
@@ -226,15 +242,6 @@ public class AmesRandomAccess extends CsvDataset {
          */
         @Override
         public AmesRandomAccess build() {
-            if (features.isEmpty()) {
-                parseFeatures();
-                for (String name : af.featureArray) {
-                    addFeature(name);
-                }
-            }
-            if (labels.isEmpty()) {
-                addNumericLabel("saleprice");
-            }
             return new AmesRandomAccess(this);
         }
 
@@ -260,6 +267,5 @@ public class AmesRandomAccess extends CsvDataset {
 
         List<String> featureArray;
         Set<String> categorical;
-        Map<String, Map<String, Integer>> featureToMap;
     }
 }
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/CsvDataset.java b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/CsvDataset.java
index 5de3eccc9..702ceaa5f 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/CsvDataset.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/CsvDataset.java
@@ -155,7 +155,6 @@ public class CsvDataset extends TabularDataset {
          * @return the new {@link CsvDataset}
          */
         public CsvDataset build() {
-            validate();
             return new CsvDataset(this);
         }
     }
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/TabularDataset.java b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/TabularDataset.java
index 43ff98dc7..9a0febfe0 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/TabularDataset.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/TabularDataset.java
@@ -44,6 +44,31 @@ public abstract class TabularDataset extends RandomAccessDataset {
         super(builder);
         features = builder.features;
         labels = builder.labels;
+
+        if (features.isEmpty()) {
+            throw new IllegalArgumentException("Missing features.");
+        }
+        if (labels.isEmpty() && !builder.allowNoLabels) {
+            throw new IllegalArgumentException("Missing labels.");
+        }
+    }
+
+    /**
+     * Gets the feature size of current {@link TabularDataset}.
+     *
+     * @return the feature size
+     */
+    public int getFeatureSize() {
+        return features.size();
+    }
+
+    /**
+     * Gets the label size of current {@link TabularDataset}.
+     *
+     * @return the feature size
+     */
+    public int getLabelSize() {
+        return labels.size();
     }
 
     /** {@inheritDoc} */
@@ -117,6 +142,7 @@ public abstract class TabularDataset extends RandomAccessDataset {
 
         protected List<Feature> features;
         protected List<Feature> labels;
+        protected boolean allowNoLabels;
 
         protected BaseBuilder() {
             features = new ArrayList<>();
@@ -265,17 +291,13 @@ public abstract class TabularDataset extends RandomAccessDataset {
         }
 
         /**
-         * Validates the builder to ensure it is correct.
+         * Indicates the dataset should not have any labels.
          *
-         * @throws IllegalArgumentException if there is an error with the builder arguments
+         * @return this builder
          */
-        protected void validate() {
-            if (features.isEmpty()) {
-                throw new IllegalArgumentException("Missing features.");
-            }
-            if (labels.isEmpty()) {
-                throw new IllegalArgumentException("Missing labels.");
-            }
+        public T noLabels() {
+            allowNoLabels = true;
+            return self();
         }
     }
 }
diff --git a/basicdataset/src/main/resources/ai/djl/basicdataset/tabular/ames.json b/basicdataset/src/main/resources/ai/djl/basicdataset/tabular/ames.json
index 07ac15318..336365f10 100644
--- a/basicdataset/src/main/resources/ai/djl/basicdataset/tabular/ames.json
+++ b/basicdataset/src/main/resources/ai/djl/basicdataset/tabular/ames.json
@@ -30,7 +30,6 @@
     "exterqual",
     "extercond",
     "foundation",
-    "bsmtqual",
     "bsmtcond",
     "bsmtexposure",
     "bsmtfintype1",
@@ -103,6 +102,7 @@
     "exterqual",
     "extercond",
     "foundation",
+    "bsmtqual",
     "bsmtcond",
     "bsmtexposure",
     "bsmtfintype1",
@@ -124,354 +124,5 @@
     "miscfeature",
     "saletype",
     "salecondition"
-  ],
-  "featureToMap": {
-    "lotconfig": {
-      "CulDSac": 3,
-      "Corner": 2,
-      "FR3": 4,
-      "FR2": 1,
-      "Inside": 0
-    },
-    "garagefinish": {
-      "Fin": 2,
-      "NA": 3,
-      "RFn": 0,
-      "Unf": 1
-    },
-    "poolqc": {
-      "NA": 0,
-      "Ex": 1,
-      "Fa": 2,
-      "Gd": 3
-    },
-    "kitchenqual": {
-      "Ex": 2,
-      "Fa": 3,
-      "Gd": 0,
-      "TA": 1
-    },
-    "functional": {
-      "Mod": 4,
-      "Min2": 3,
-      "Sev": 6,
-      "Maj2": 5,
-      "Typ": 0,
-      "Min1": 1,
-      "Maj1": 2
-    },
-    "mszoning": {
-      "FV": 3,
-      "RH": 4,
-      "RL": 0,
-      "RM": 1,
-      "C (all)": 2
-    },
-    "bsmtexposure": {
-      "No": 0,
-      "Mn": 2,
-      "NA": 4,
-      "Av": 3,
-      "Gd": 1
-    },
-    "heatingqc": {
-      "Ex": 0,
-      "Fa": 3,
-      "Gd": 1,
-      "TA": 2,
-      "Po": 4
-    },
-    "foundation": {
-      "CBlock": 1,
-      "BrkTil": 2,
-      "Wood": 3,
-      "PConc": 0,
-      "Slab": 4,
-      "Stone": 5
-    },
-    "utilities": {
-      "NoSeWa": 1,
-      "AllPub": 0
-    },
-    "paveddrive": {
-      "P": 2,
-      "Y": 0,
-      "N": 1
-    },
-    "exterior1st": {
-      "AsbShng": 8,
-      "Stucco": 9,
-      "BrkComm": 10,
-      "AsphShn": 11,
-      "CBlock": 14,
-      "WdShing": 5,
-      "Stone": 12,
-      "BrkFace": 4,
-      "MetalSd": 1,
-      "Wd Sdng": 2,
-      "HdBoard": 3,
-      "VinylSd": 0,
-      "Plywood": 7,
-      "CemntBd": 6,
-      "ImStucc": 13
-    },
-    "roofstyle": {
-      "Mansard": 3,
-      "Shed": 5,
-      "Gambrel": 2,
-      "Gable": 0,
-      "Flat": 4,
-      "Hip": 1
-    },
-    "street": {
-      "Pave": 0,
-      "Grvl": 1
-    },
-    "alley": {
-      "NA": 0,
-      "Grvl": 1,
-      "Pave": 2
-    },
-    "heating": {
-      "Grav": 2,
-      "OthW": 4,
-      "Floor": 5,
-      "Wall": 3,
-      "GasA": 0,
-      "GasW": 1
-    },
-    "bldgtype": {
-      "Twnhs": 4,
-      "1Fam": 0,
-      "TwnhsE": 3,
-      "2fmCon": 1,
-      "Duplex": 2
-    },
-    "masvnrtype": {
-      "BrkFace": 0,
-      "NA": 4,
-      "None": 1,
-      "Stone": 2,
-      "BrkCmn": 3
-    },
-    "extercond": {
-      "Ex": 4,
-      "Fa": 2,
-      "TA": 0,
-      "Gd": 1,
-      "Po": 3
-    },
-    "garagetype": {
-      "CarPort": 3,
-      "Detchd": 1,
-      "NA": 4,
-      "Basment": 5,
-      "Attchd": 0,
-      "2Types": 6,
-      "BuiltIn": 2
-    },
-    "electrical": {
-      "FuseA": 2,
-      "NA": 5,
-      "FuseF": 1,
-      "SBrkr": 0,
-      "Mix": 4,
-      "FuseP": 3
-    },
-    "exterqual": {
-      "Ex": 2,
-      "Fa": 3,
-      "Gd": 0,
-      "TA": 1
-    },
-    "landslope": {
-      "Mod": 1,
-      "Sev": 2,
-      "Gtl": 0
-    },
-    "bsmtcond": {
-      "NA": 2,
-      "Fa": 3,
-      "TA": 0,
-      "Gd": 1,
-      "Po": 4
-    },
-    "landcontour": {
-      "Lvl": 0,
-      "Low": 2,
-      "Bnk": 1,
-      "HLS": 3
-    },
-    "fireplacequ": {
-      "NA": 0,
-      "Ex": 4,
-      "Fa": 3,
-      "TA": 1,
-      "Gd": 2,
-      "Po": 5
-    },
-    "housestyle": {
-      "1Story": 1,
-      "2.5Fin": 7,
-      "2.5Unf": 6,
-      "2Story": 0,
-      "SLvl": 5,
-      "1.5Unf": 3,
-      "SFoyer": 4,
-      "1.5Fin": 2
-    },
-    "exterior2nd": {
-      "Stucco": 8,
-      "AsbShng": 9,
-      "AsphShn": 12,
-      "CBlock": 15,
-      "CmentBd": 6,
-      "Brk Cmn": 10,
-      "Wd Shng": 2,
-      "Stone": 13,
-      "BrkFace": 7,
-      "MetalSd": 1,
-      "HdBoard": 3,
-      "Wd Sdng": 5,
-      "VinylSd": 0,
-      "Plywood": 4,
-      "ImStucc": 11,
-      "Other": 14
-    },
-    "condition1": {
-      "RRAe": 4,
-      "RRNn": 5,
-      "RRAn": 6,
-      "PosA": 7,
-      "Artery": 3,
-      "RRNe": 8,
-      "Norm": 0,
-      "PosN": 2,
-      "Feedr": 1
-    },
-    "garagequal": {
-      "NA": 3,
-      "Ex": 4,
-      "Fa": 1,
-      "TA": 0,
-      "Gd": 2,
-      "Po": 5
-    },
-    "roofmatl": {
-      "Tar\u0026Grv": 5,
-      "ClyTile": 7,
-      "WdShngl": 1,
-      "Membran": 4,
-      "Metal": 2,
-      "Roll": 6,
-      "CompShg": 0,
-      "WdShake": 3
-    },
-    "condition2": {
-      "RRAe": 7,
-      "RRNn": 2,
-      "Norm": 0,
-      "PosA": 5,
-      "RRAn": 6,
-      "PosN": 4,
-      "Feedr": 3,
-      "Artery": 1
-    },
-    "saletype": {
-      "New": 1,
-      "Oth": 8,
-      "Con": 7,
-      "ConLw": 6,
-      "WD": 0,
-      "CWD": 5,
-      "COD": 2,
-      "ConLI": 4,
-      "ConLD": 3
-    },
-    "garagecond": {
-      "NA": 2,
-      "Ex": 5,
-      "Fa": 1,
-      "TA": 0,
-      "Gd": 3,
-      "Po": 4
-    },
-    "centralair": {
-      "Y": 0,
-      "N": 1
-    },
-    "lotshape": {
-      "IR2": 2,
-      "IR1": 1,
-      "IR3": 3,
-      "Reg": 0
-    },
-    "neighborhood": {
-      "NWAmes": 6,
-      "BrDale": 22,
-      "Somerst": 5,
-      "Timber": 16,
-      "Gilbert": 17,
-      "SawyerW": 12,
-      "NPkVill": 20,
-      "Sawyer": 9,
-      "NoRidge": 3,
-      "NridgHt": 10,
-      "StoneBr": 18,
-      "Veenker": 1,
-      "BrkSide": 8,
-      "NAmes": 11,
-      "CollgCr": 0,
-      "MeadowV": 14,
-      "Blmngtn": 21,
-      "Blueste": 24,
-      "IDOTRR": 13,
-      "Edwards": 15,
-      "ClearCr": 19,
-      "Mitchel": 4,
-      "OldTown": 7,
-      "Crawfor": 2,
-      "SWISU": 23
-    },
-    "fence": {
-      "GdPrv": 3,
-      "NA": 0,
-      "MnPrv": 1,
-      "GdWo": 2,
-      "MnWw": 4
-    },
-    "bsmtfintype2": {
-      "Rec": 4,
-      "NA": 2,
-      "BLQ": 1,
-      "ALQ": 3,
-      "LwQ": 5,
-      "Unf": 0,
-      "GLQ": 6
-    },
-    "miscfeature": {
-      "TenC": 4,
-      "NA": 0,
-      "Shed": 1,
-      "Othr": 3,
-      "Gar2": 2
-    },
-    "salecondition": {
-      "AdjLand": 3,
-      "Alloca": 4,
-      "Partial": 2,
-      "Family": 5,
-      "Normal": 0,
-      "Abnorml": 1
-    },
-    "bsmtfintype1": {
-      "Rec": 3,
-      "NA": 5,
-      "ALQ": 1,
-      "BLQ": 4,
-      "LwQ": 6,
-      "Unf": 2,
-      "GLQ": 0
-    }
-  }
+  ]
 }
diff --git a/djl-zero/src/main/java/ai/djl/zero/cv/ObjectDetection.java b/djl-zero/src/main/java/ai/djl/zero/cv/ObjectDetection.java
new file mode 100644
index 000000000..a10727f65
--- /dev/null
+++ b/djl-zero/src/main/java/ai/djl/zero/cv/ObjectDetection.java
@@ -0,0 +1,137 @@
+/*
+ * 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.zero.cv;
+
+import ai.djl.Model;
+import ai.djl.basicdataset.cv.ObjectDetectionDataset;
+import ai.djl.basicmodelzoo.cv.object_detection.ssd.SingleShotDetection;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.transform.ToTensor;
+import ai.djl.modality.cv.translator.SingleShotDetectionTranslator;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Block;
+import ai.djl.nn.SequentialBlock;
+import ai.djl.repository.zoo.ZooModel;
+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.Dataset;
+import ai.djl.training.evaluator.BoundingBoxError;
+import ai.djl.training.evaluator.SingleShotDetectionAccuracy;
+import ai.djl.training.listener.TrainingListener;
+import ai.djl.training.loss.SingleShotDetectionLoss;
+import ai.djl.translate.TranslateException;
+import ai.djl.translate.Translator;
+import ai.djl.zero.Performance;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+/** ObjectDetection takes an image and extract one or more main subjects in the image. */
+public final class ObjectDetection {
+    private ObjectDetection() {}
+
+    /**
+     * Trains the recommended object detection model on a custom dataset. Currently, trains a
+     * SingleShotDetection Model.
+     *
+     * <p>In order to train on a custom dataset, you must create a custom {@link
+     * ObjectDetectionDataset} to load your data.
+     *
+     * @param dataset the data to train with
+     * @param performance to determine the desired model tradeoffs
+     * @return the model as a {@link ZooModel} with the {@link Translator} included
+     * @throws IOException if the dataset could not be loaded
+     * @throws TranslateException if the translator has errors
+     */
+    public static ZooModel<Image, DetectedObjects> train(
+            ObjectDetectionDataset dataset, Performance performance)
+            throws IOException, TranslateException {
+        List<String> classes = dataset.getClasses();
+        int channels = dataset.getImageChannels();
+        int width =
+                dataset.getImageWidth()
+                        .orElseThrow(
+                                () ->
+                                        new IllegalArgumentException(
+                                                "The dataset must have a fixed image width"));
+        int height =
+                dataset.getImageHeight()
+                        .orElseThrow(
+                                () ->
+                                        new IllegalArgumentException(
+                                                "The dataset must have a fixed image height"));
+
+        Shape imageShape = new Shape(channels, height, width);
+
+        Dataset[] splitDataset = dataset.randomSplit(8, 2);
+        Dataset trainDataset = splitDataset[0];
+        Dataset validateDataset = splitDataset[1];
+
+        Block block = getSsdTrainBlock(classes.size());
+        Model model = Model.newInstance("ObjectDetection");
+        model.setBlock(block);
+
+        TrainingConfig trainingConfig =
+                new DefaultTrainingConfig(new SingleShotDetectionLoss())
+                        .addEvaluator(new SingleShotDetectionAccuracy("classAccuracy"))
+                        .addEvaluator(new BoundingBoxError("boundingBoxError"))
+                        .addTrainingListeners(TrainingListener.Defaults.basic());
+
+        try (Trainer trainer = model.newTrainer(trainingConfig)) {
+            trainer.initialize(new Shape(1).addAll(imageShape));
+            EasyTrain.fit(trainer, 50, trainDataset, validateDataset);
+        }
+
+        Translator<Image, DetectedObjects> translator =
+                SingleShotDetectionTranslator.builder()
+                        .addTransform(new ToTensor())
+                        .optSynset(classes)
+                        .optThreshold(0.6f)
+                        .build();
+
+        return new ZooModel<>(model, translator);
+    }
+
+    private static Block getSsdTrainBlock(int numClasses) {
+        int[] numFilters = {16, 32, 64};
+        SequentialBlock baseBlock = new SequentialBlock();
+        for (int numFilter : numFilters) {
+            baseBlock.add(SingleShotDetection.getDownSamplingBlock(numFilter));
+        }
+
+        List<List<Float>> sizes = new ArrayList<>();
+        List<List<Float>> ratios = new ArrayList<>();
+        for (int i = 0; i < 5; i++) {
+            ratios.add(Arrays.asList(1f, 2f, 0.5f));
+        }
+        sizes.add(Arrays.asList(0.2f, 0.272f));
+        sizes.add(Arrays.asList(0.37f, 0.447f));
+        sizes.add(Arrays.asList(0.54f, 0.619f));
+        sizes.add(Arrays.asList(0.71f, 0.79f));
+        sizes.add(Arrays.asList(0.88f, 0.961f));
+
+        return SingleShotDetection.builder()
+                .setNumClasses(numClasses)
+                .setNumFeatures(3)
+                .optGlobalPool(true)
+                .setRatios(ratios)
+                .setSizes(sizes)
+                .setBaseNetwork(baseBlock)
+                .build();
+    }
+}
diff --git a/djl-zero/src/main/java/ai/djl/zero/tabular/Tabular.java b/djl-zero/src/main/java/ai/djl/zero/tabular/Tabular.java
new file mode 100644
index 000000000..634466a0d
--- /dev/null
+++ b/djl-zero/src/main/java/ai/djl/zero/tabular/Tabular.java
@@ -0,0 +1,97 @@
+/*
+ * 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.zero.tabular;
+
+import ai.djl.Model;
+import ai.djl.basicdataset.tabular.TabularDataset;
+import ai.djl.basicmodelzoo.tabular.TabNet;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Block;
+import ai.djl.repository.zoo.ZooModel;
+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.Dataset;
+import ai.djl.training.listener.TrainingListener;
+import ai.djl.training.loss.TabNetRegressionLoss;
+import ai.djl.translate.NoopTranslator;
+import ai.djl.translate.TranslateException;
+import ai.djl.zero.Performance;
+
+import java.io.IOException;
+
+/** Tabular takes a NDList as input and output an NDList (for supervised learning). */
+public final class Tabular {
+
+    private Tabular() {}
+
+    /**
+     * Trains a Model on a custom dataset. Currently, trains a TabNet Model.
+     *
+     * <p>In order to train on a custom dataset, you must create a custom {@link TabularDataset} to
+     * load your data.
+     *
+     * @param dataset the data to train with
+     * @param performance to determine the desired model tradeoffs
+     * @return the model as a {@link ZooModel}
+     * @throws IOException if the dataset could not be loaded
+     * @throws TranslateException if the translator has errors
+     */
+    public static ZooModel<NDList, NDList> train(TabularDataset dataset, Performance performance)
+            throws IOException, TranslateException {
+        Dataset[] splitDataset = dataset.randomSplit(8, 2);
+        Dataset trainDataset = splitDataset[0];
+        Dataset validateDataset = splitDataset[1];
+        int featureSize = dataset.getFeatureSize();
+        int labelSize = dataset.getLabelSize();
+
+        Block block;
+        if (performance.equals(Performance.FAST)) {
+            // for fast cases, we set the number of independent layers and shared layers lower
+            block =
+                    TabNet.builder()
+                            .setInputDim(featureSize)
+                            .setOutDim(labelSize)
+                            .optNumIndependent(1)
+                            .optNumShared(1)
+                            .build();
+        } else if (performance.equals(Performance.BALANCED)) {
+            block = TabNet.builder().setInputDim(featureSize).setOutDim(labelSize).build();
+        } else {
+            // for accurate cases, we set the number of independent layers and shared layers higher
+            block =
+                    TabNet.builder()
+                            .setInputDim(featureSize)
+                            .setOutDim(labelSize)
+                            .optNumIndependent(4)
+                            .optNumShared(4)
+                            .build();
+        }
+
+        Model model = Model.newInstance("tabular");
+        model.setBlock(block);
+
+        TrainingConfig trainingConfig =
+                new DefaultTrainingConfig(new TabNetRegressionLoss())
+                        .addTrainingListeners(TrainingListener.Defaults.basic());
+
+        try (Trainer trainer = model.newTrainer(trainingConfig)) {
+            trainer.initialize(new Shape(1, featureSize));
+            EasyTrain.fit(trainer, 20, trainDataset, validateDataset);
+        }
+
+        return new ZooModel<>(model, new NoopTranslator());
+    }
+}
diff --git a/djl-zero/src/main/java/ai/djl/zero/tabular/package-info.java b/djl-zero/src/main/java/ai/djl/zero/tabular/package-info.java
new file mode 100644
index 000000000..42c8aebc1
--- /dev/null
+++ b/djl-zero/src/main/java/ai/djl/zero/tabular/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * 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.
+ */
+
+/** Contains training for Tabular({@link ai.djl.Application.Tabular}) tasks. */
+package ai.djl.zero.tabular;
diff --git a/engines/mxnet/mxnet-engine/README.md b/engines/mxnet/mxnet-engine/README.md
index 3c72f3e9d..237623d41 100644
--- a/engines/mxnet/mxnet-engine/README.md
+++ b/engines/mxnet/mxnet-engine/README.md
@@ -73,11 +73,11 @@ installed on your GPU machine, you can use one of the following library:
 **Important:** Since cuda 11.0, you must install CUDNN and NCCL matches the CUDA version manually.
 Apache MXNet no longer statically link CUDNN and NCCL in it.
 
-Apache MXNet 1.8.0 cu110 package supports SM 5.0, 6.0, 7.0 and 8.0 cuda architectures.
-Apache MXNet 1.8.0 cu102 package supports SM 3.0, 5.0, 6.0, 7.0 and 7.5 cuda architectures.
+Apache MXNet 1.9.1 cu112 package supports SM 5.0, 6.0, 7.0 and 8.0 cuda architectures.
+Apache MXNet 1.9.1 cu102 package supports SM 3.0, 5.0, 6.0, 7.0 and 7.5 cuda architectures.
 
-- ai.djl.mxnet:mxnet-native-cu102mkl:1.9.1:linux-x86_64 - CUDA 11.2
-- ai.djl.mxnet:mxnet-native-cu101mkl:1.9.1:linux-x86_64 - CUDA 10.2
+- ai.djl.mxnet:mxnet-native-cu112mkl:1.9.1:linux-x86_64 - CUDA 11.2
+- ai.djl.mxnet:mxnet-native-cu102mkl:1.9.1:linux-x86_64 - CUDA 10.2
 
 ```xml
 <dependency>
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
index 5ae9e0494..192b1c444 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
@@ -1544,6 +1544,10 @@ public final class JniUtils {
             String[] extraFileKeys,
             String[] extraFileValues) {
         Device device = manager.getDevice();
+        // MPS doesn't support mapLocation
+        if ("mps".equals(device.getDeviceType())) {
+            mapLocation = false;
+        }
         logger.debug("mapLocation: {}", mapLocation);
         logger.debug("extraFileKeys: {}", Arrays.toString(extraFileKeys));
         long handle =
@@ -1571,6 +1575,11 @@ public final class JniUtils {
         if (hasSize) {
             size = new DataInputStream(is).readLong();
         }
+        // MPS doesn't support mapLocation
+        if ("mps".equals(device.getDeviceType())) {
+            mapLocation = false;
+        }
+        logger.debug("mapLocation: {}", mapLocation);
         return PyTorchLibrary.LIB.moduleLoad(
                 is,
                 new int[] {PtDeviceType.toDeviceType(device), device.getDeviceId()},
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java
index 331e1ff52..3eec14fbc 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java
@@ -14,17 +14,31 @@
 package ai.djl.tensorflow.engine;
 
 import org.tensorflow.internal.c_api.TF_Graph;
+import org.tensorflow.internal.c_api.TF_Operation;
 import org.tensorflow.internal.c_api.TF_Session;
+import org.tensorflow.internal.c_api.global.tensorflow;
+import org.tensorflow.proto.framework.CollectionDef;
 import org.tensorflow.proto.framework.MetaGraphDef;
+import org.tensorflow.proto.framework.SignatureDef;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /** The wrapper class for native resources required for SavedModelBundle. */
 public class SavedModelBundle implements AutoCloseable {
 
+    private static final String INIT_OP_SIGNATURE_KEY = "__saved_model_init_op";
+    private static final String MAIN_OP_COLLECTION_KEY = "saved_model_main_op";
+    private static final String LEGACY_INIT_OP_COLLECTION_KEY = "legacy_init_op";
+    private static final String TABLE_INITIALIZERS_COLLECTION_KEY = "table_initializer";
+
     private TF_Graph graphHandle;
     private TF_Session sessionHandle;
     private MetaGraphDef metaGraphDef;
+    private TF_Operation[] targetOpHandles;
     private AtomicBoolean closed;
 
     public SavedModelBundle(
@@ -33,6 +47,61 @@ public class SavedModelBundle implements AutoCloseable {
         this.sessionHandle = sessionHandle;
         this.metaGraphDef = metaGraphDef;
         closed = new AtomicBoolean(false);
+
+        Map<String, SignatureDef> functions = new ConcurrentHashMap<>();
+        metaGraphDef
+                .getSignatureDefMap()
+                .forEach(
+                        (signatureName, signatureDef) -> {
+                            if (!functions.containsKey(signatureName)) {
+                                functions.put(signatureName, signatureDef);
+                            }
+                        });
+
+        List<TF_Operation> initOps = new ArrayList<>();
+        TF_Operation initOp = findInitOp(functions, metaGraphDef.getCollectionDefMap());
+        if (initOp != null) {
+            initOps.add(initOp);
+        }
+
+        if (metaGraphDef.containsCollectionDef(TABLE_INITIALIZERS_COLLECTION_KEY)) {
+            metaGraphDef
+                    .getCollectionDefMap()
+                    .get(TABLE_INITIALIZERS_COLLECTION_KEY)
+                    .getNodeList()
+                    .getValueList()
+                    .forEach(
+                            node -> {
+                                initOps.add(tensorflow.TF_GraphOperationByName(graphHandle, node));
+                            });
+        }
+        targetOpHandles = initOps.toArray(new TF_Operation[0]);
+    }
+
+    private TF_Operation findInitOp(
+            Map<String, SignatureDef> signatures, Map<String, CollectionDef> collections) {
+        SignatureDef initSig = signatures.get(INIT_OP_SIGNATURE_KEY);
+        if (initSig != null) {
+            String opName = initSig.getOutputsMap().get(INIT_OP_SIGNATURE_KEY).getName();
+            return tensorflow.TF_GraphOperationByName(graphHandle, opName);
+        }
+
+        CollectionDef initCollection;
+        if (collections.containsKey(MAIN_OP_COLLECTION_KEY)) {
+            initCollection = collections.get(MAIN_OP_COLLECTION_KEY);
+        } else {
+            initCollection = collections.get(LEGACY_INIT_OP_COLLECTION_KEY);
+        }
+
+        if (initCollection != null) {
+            CollectionDef.NodeList nodes = initCollection.getNodeList();
+            if (nodes.getValueCount() != 1) {
+                throw new IllegalArgumentException("Expected exactly one main op in saved model.");
+            }
+            String opName = nodes.getValue(0);
+            return tensorflow.TF_GraphOperationByName(graphHandle, opName);
+        }
+        return null;
     }
 
     /**
@@ -62,6 +131,10 @@ public class SavedModelBundle implements AutoCloseable {
         return metaGraphDef;
     }
 
+    TF_Operation[] getTargetOpHandles() {
+        return targetOpHandles;
+    }
+
     /** {@inheritDoc} */
     @Override
     public void close() {
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java
index b6a853bdc..2453e58c2 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java
@@ -67,6 +67,7 @@ public class TfSymbolBlock extends AbstractSymbolBlock implements AutoCloseable
         this.bundle = bundle;
         graphHandle = bundle.getGraph();
         sessionHandle = bundle.getSession();
+        targetOpHandles = bundle.getTargetOpHandles();
         MetaGraphDef metaGraphDef = bundle.getMetaGraphDef();
         Map<String, SignatureDef> signatureDefMap = metaGraphDef.getSignatureDefMap();
         if (signatureDefMap.containsKey(signatureDefKey)) {
@@ -89,8 +90,6 @@ public class TfSymbolBlock extends AbstractSymbolBlock implements AutoCloseable
         }
         describeInput();
         describeOutput();
-        // we don't use target for now
-        targetOpHandles = new TF_Operation[0];
     }
 
     /** {@inheritDoc} */
@@ -215,10 +214,6 @@ public class TfSymbolBlock extends AbstractSymbolBlock implements AutoCloseable
             for (String key : keys) {
                 TensorInfo tensorInfo = outputsMap.get(key);
                 TensorShapeProto shapeProto = tensorInfo.getTensorShape();
-                // does not support string tensors
-                if (tensorInfo.getDtype() == org.tensorflow.proto.framework.DataType.DT_STRING) {
-                    continue;
-                }
                 outputDescriptions.add(
                         key,
                         new Shape(
diff --git a/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java b/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java
index 92a677484..ef281e2c0 100644
--- a/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java
+++ b/examples/src/main/java/ai/djl/examples/inference/SemanticSegmentation.java
@@ -12,14 +12,12 @@
  */
 package ai.djl.examples.inference;
 
-import ai.djl.Application;
 import ai.djl.ModelException;
 import ai.djl.inference.Predictor;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
-import ai.djl.modality.cv.transform.Normalize;
-import ai.djl.modality.cv.transform.ToTensor;
-import ai.djl.modality.cv.translator.SemanticSegmentationTranslator;
+import ai.djl.modality.cv.output.CategoryMask;
+import ai.djl.modality.cv.translator.SemanticSegmentationTranslatorFactory;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.training.util.ProgressBar;
@@ -28,6 +26,7 @@ import ai.djl.translate.TranslateException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.awt.Color;
 import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
@@ -42,8 +41,7 @@ import java.nio.file.Paths;
  */
 public final class SemanticSegmentation {
 
-    private static final Logger logger =
-            LoggerFactory.getLogger(SemanticSegmentationTranslator.class);
+    private static final Logger logger = LoggerFactory.getLogger(SemanticSegmentation.class);
 
     private SemanticSegmentation() {}
 
@@ -53,41 +51,58 @@ public final class SemanticSegmentation {
 
     public static void predict() throws IOException, ModelException, TranslateException {
         Path imageFile = Paths.get("src/test/resources/dog_bike_car.jpg");
-        Image img = ImageFactory.getInstance().fromFile(imageFile);
-
-        float[] mean = {0.485f, 0.456f, 0.406f};
-        float[] std = {0.229f, 0.224f, 0.225f};
+        ImageFactory factory = ImageFactory.getInstance();
+        Image img = factory.fromFile(imageFile);
 
         String url =
                 "https://mlrepo.djl.ai/model/cv/semantic_segmentation/ai/djl/pytorch/deeplabv3/0.0.1/deeplabv3.zip";
-        SemanticSegmentationTranslator translator =
-                SemanticSegmentationTranslator.builder()
-                        .addTransform(new ToTensor())
-                        .addTransform(new Normalize(mean, std))
-                        .build();
 
-        Criteria<Image, Image> criteria =
+        Criteria<Image, CategoryMask> criteria =
                 Criteria.builder()
-                        .optApplication(Application.CV.SEMANTIC_SEGMENTATION)
-                        .setTypes(Image.class, Image.class)
+                        .setTypes(Image.class, CategoryMask.class)
                         .optModelUrls(url)
-                        .optTranslator(translator)
+                        .optTranslatorFactory(new SemanticSegmentationTranslatorFactory())
                         .optEngine("PyTorch")
                         .optProgress(new ProgressBar())
                         .build();
-        try (ZooModel<Image, Image> model = criteria.loadModel();
-                Predictor<Image, Image> predictor = model.newPredictor()) {
-            Image semanticImage = predictor.predict(img);
-            saveSemanticImage(semanticImage);
+        Image bg = factory.fromFile(Paths.get("../website/img/stars-in-the-night-sky.jpg"));
+        try (ZooModel<Image, CategoryMask> model = criteria.loadModel();
+                Predictor<Image, CategoryMask> predictor = model.newPredictor()) {
+            CategoryMask mask = predictor.predict(img);
+
+            // Highlights the detected object on the image with random opaque colors.
+            Image img1 = img.duplicate();
+            mask.drawMask(img1, 255);
+            saveSemanticImage(img1, "semantic_instances1.png");
+
+            // Highlights the detected object on the image with random colors.
+            Image img2 = img.duplicate();
+            mask.drawMask(img2, 180, 0);
+            saveSemanticImage(img2, "semantic_instances2.png");
+
+            // Highlights only the dog with blue color.
+            Image img3 = img.duplicate();
+            mask.drawMask(img3, 12, Color.BLUE.brighter().getRGB(), 180);
+            saveSemanticImage(img3, "semantic_instances3.png");
+
+            // Extract dog from the image
+            Image dog = mask.getMaskImage(img, 12);
+            dog = dog.resize(img.getWidth(), img.getHeight(), true);
+            saveSemanticImage(dog, "semantic_instances4.png");
+
+            // Replace background with an image
+            bg = bg.resize(img.getWidth(), img.getHeight(), true);
+            bg.drawImage(dog, true);
+            saveSemanticImage(bg, "semantic_instances5.png");
         }
     }
 
-    private static void saveSemanticImage(Image img) throws IOException {
+    private static void saveSemanticImage(Image img, String fileName) throws IOException {
         Path outputDir = Paths.get("build/output");
         Files.createDirectories(outputDir);
 
-        Path imagePath = outputDir.resolve("semantic_instances.png");
+        Path imagePath = outputDir.resolve(fileName);
         img.save(Files.newOutputStream(imagePath), "png");
-        logger.info("Segmentation result image has been saved in: {}", imagePath);
+        logger.info("Segmentation result image has been saved in: {}", imagePath.toAbsolutePath());
     }
 }
diff --git a/examples/src/main/java/ai/djl/examples/training/TrainAirfoilWithTabNet.java b/examples/src/main/java/ai/djl/examples/training/TrainAirfoilWithTabNet.java
new file mode 100644
index 000000000..77e3fae0c
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/training/TrainAirfoilWithTabNet.java
@@ -0,0 +1,111 @@
+/*
+ * 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.examples.training;
+
+import ai.djl.Model;
+import ai.djl.basicdataset.tabular.AirfoilRandomAccess;
+import ai.djl.basicmodelzoo.tabular.TabNet;
+import ai.djl.engine.Engine;
+import ai.djl.examples.training.util.Arguments;
+import ai.djl.metric.Metrics;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Block;
+import ai.djl.training.DefaultTrainingConfig;
+import ai.djl.training.EasyTrain;
+import ai.djl.training.Trainer;
+import ai.djl.training.TrainingResult;
+import ai.djl.training.dataset.Dataset;
+import ai.djl.training.dataset.RandomAccessDataset;
+import ai.djl.training.listener.SaveModelTrainingListener;
+import ai.djl.training.listener.TrainingListener;
+import ai.djl.training.loss.TabNetRegressionLoss;
+import ai.djl.training.util.ProgressBar;
+import ai.djl.translate.TranslateException;
+
+import java.io.IOException;
+
+public final class TrainAirfoilWithTabNet {
+    private TrainAirfoilWithTabNet() {}
+
+    public static void main(String[] args) throws TranslateException, IOException {
+        TrainAirfoilWithTabNet.runExample(args);
+    }
+
+    public static TrainingResult runExample(String[] args) throws IOException, TranslateException {
+        Arguments arguments = new Arguments().parseArgs(args);
+        if (arguments == null) {
+            return null;
+        }
+
+        // Construct a tabNet instance
+        Block tabNet = TabNet.builder().setInputDim(5).setOutDim(1).build();
+
+        try (Model model = Model.newInstance("tabNet")) {
+            model.setBlock(tabNet);
+
+            // get the training and validation dataset
+            RandomAccessDataset[] randomAccessDatasets = getDataset(arguments);
+            RandomAccessDataset trainingSet = randomAccessDatasets[0];
+            RandomAccessDataset validateSet = randomAccessDatasets[1];
+
+            // setup training configuration
+            DefaultTrainingConfig config = setupTrainingConfig(arguments);
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                trainer.setMetrics(new Metrics());
+
+                Shape inputShape = new Shape(arguments.getBatchSize(), 5);
+
+                // initialize trainer with proper input shape
+                trainer.initialize(inputShape);
+
+                EasyTrain.fit(trainer, arguments.getEpoch(), trainingSet, validateSet);
+
+                return trainer.getTrainingResult();
+            }
+        }
+    }
+
+    private static DefaultTrainingConfig setupTrainingConfig(Arguments arguments) {
+        String outputDir = arguments.getOutputDir();
+        SaveModelTrainingListener listener = new SaveModelTrainingListener(outputDir);
+        listener.setSaveModelCallback(
+                trainer -> {
+                    TrainingResult result = trainer.getTrainingResult();
+                    Model model = trainer.getModel();
+                    model.setProperty("Loss", String.format("%.5f", result.getValidateLoss()));
+                });
+
+        return new DefaultTrainingConfig(new TabNetRegressionLoss())
+                .optDevices(Engine.getInstance().getDevices(arguments.getMaxGpus()))
+                .addTrainingListeners(TrainingListener.Defaults.logging(outputDir))
+                .addTrainingListeners(listener);
+    }
+
+    private static RandomAccessDataset[] getDataset(Arguments arguments)
+            throws IOException, TranslateException {
+        AirfoilRandomAccess.Builder airfoilBuilder = AirfoilRandomAccess.builder();
+
+        // only train dataset is available, so we get train dataset and split them
+        airfoilBuilder.optUsage(Dataset.Usage.TRAIN).setSampling(arguments.getBatchSize(), true);
+
+        for (int i = 0; i < airfoilBuilder.getAvailableFeatures().size() - 1; i++) {
+            airfoilBuilder.addFeature(airfoilBuilder.getAvailableFeatures().get(i));
+        }
+
+        AirfoilRandomAccess airfoilRandomAccess = airfoilBuilder.build();
+        airfoilRandomAccess.prepare(new ProgressBar());
+        // split the dataset into
+        return airfoilRandomAccess.randomSplit(8, 2);
+    }
+}
diff --git a/examples/src/main/java/ai/djl/examples/training/TrainPikachu.java b/examples/src/main/java/ai/djl/examples/training/TrainPikachu.java
index 1c11df399..534f58d85 100644
--- a/examples/src/main/java/ai/djl/examples/training/TrainPikachu.java
+++ b/examples/src/main/java/ai/djl/examples/training/TrainPikachu.java
@@ -134,7 +134,7 @@ public final class TrainPikachu {
                         .optUsage(usage)
                         .optLimit(arguments.getLimit())
                         .optPipeline(pipeline)
-                        .setSampling(arguments.getBatchSize(), true)
+                        .setSampling(1, true)
                         .build();
         pikachuDetection.prepare(new ProgressBar());
 
diff --git a/examples/src/main/java/ai/djl/examples/training/TrainPikachuWithYOLOV3.java b/examples/src/main/java/ai/djl/examples/training/TrainPikachuWithYOLOV3.java
new file mode 100644
index 000000000..adc72e78e
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/training/TrainPikachuWithYOLOV3.java
@@ -0,0 +1,113 @@
+/*
+ * 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.examples.training;
+
+import ai.djl.MalformedModelException;
+import ai.djl.Model;
+import ai.djl.basicdataset.cv.PikachuDetection;
+import ai.djl.basicmodelzoo.cv.object_detection.yolo.YOLOV3;
+import ai.djl.engine.Engine;
+import ai.djl.examples.training.util.Arguments;
+import ai.djl.metric.Metrics;
+import ai.djl.modality.cv.transform.ToTensor;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.training.DefaultTrainingConfig;
+import ai.djl.training.EasyTrain;
+import ai.djl.training.Trainer;
+import ai.djl.training.TrainingResult;
+import ai.djl.training.dataset.Dataset;
+import ai.djl.training.dataset.RandomAccessDataset;
+import ai.djl.training.listener.SaveModelTrainingListener;
+import ai.djl.training.listener.TrainingListener;
+import ai.djl.training.loss.YOLOv3Loss;
+import ai.djl.training.util.ProgressBar;
+import ai.djl.translate.Pipeline;
+import ai.djl.translate.TranslateException;
+
+import java.io.IOException;
+
+public final class TrainPikachuWithYOLOV3 {
+    private TrainPikachuWithYOLOV3() {}
+
+    public static void main(String[] args)
+            throws IOException, TranslateException, MalformedModelException {
+        TrainPikachuWithYOLOV3.runExample(args);
+    }
+
+    public static TrainingResult runExample(String[] args) throws IOException, TranslateException {
+        Arguments arguments = new Arguments().parseArgs(args);
+        if (arguments == null) {
+            return null;
+        }
+
+        try (Model model = Model.newInstance("pikachu-yolov3")) {
+            model.setBlock(YOLOV3.builder().setNumClasses(1).build());
+            RandomAccessDataset trainingSet = getDataset(Dataset.Usage.TRAIN, arguments);
+            RandomAccessDataset validateSet = getDataset(Dataset.Usage.TEST, arguments);
+
+            DefaultTrainingConfig config = setupTrainingConfig(arguments);
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                trainer.setMetrics(new Metrics());
+
+                Shape inputShape = new Shape(1, 3, 256, 256);
+                trainer.initialize(inputShape);
+
+                EasyTrain.fit(trainer, arguments.getEpoch(), trainingSet, validateSet);
+
+                return trainer.getTrainingResult();
+            }
+        }
+    }
+
+    private static RandomAccessDataset getDataset(Dataset.Usage usage, Arguments arguments)
+            throws IOException {
+        Pipeline pipeline = new Pipeline(new ToTensor());
+        PikachuDetection pikachuDetection =
+                PikachuDetection.builder()
+                        .optUsage(usage)
+                        .optLimit(arguments.getLimit())
+                        .optPipeline(pipeline)
+                        .setSampling(8, true)
+                        .build();
+        pikachuDetection.prepare(new ProgressBar());
+
+        return pikachuDetection;
+    }
+
+    private static DefaultTrainingConfig setupTrainingConfig(Arguments arguments) {
+        String outputDir = arguments.getOutputDir();
+        SaveModelTrainingListener listener = new SaveModelTrainingListener(outputDir);
+        listener.setSaveModelCallback(
+                trainer -> {
+                    TrainingResult result = trainer.getTrainingResult();
+                    Model model = trainer.getModel();
+                    model.setProperty("Loss", String.format("%.5f", result.getValidateLoss()));
+                });
+
+        float[] anchorsArray = YOLOv3Loss.getPresetAnchors();
+        for (int i = 0; i < anchorsArray.length; i++) {
+            anchorsArray[i] = anchorsArray[i] * 256 / 416; // reshaping into the
+        }
+
+        return new DefaultTrainingConfig(
+                        YOLOv3Loss.builder()
+                                .setNumClasses(1)
+                                .setInputShape(new Shape(256, 256))
+                                .setAnchorsArray(anchorsArray)
+                                .build())
+                .optDevices(Engine.getInstance().getDevices(arguments.getMaxGpus()))
+                .addTrainingListeners(TrainingListener.Defaults.basic())
+                .addTrainingListeners(listener);
+    }
+}
diff --git a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
index 9aac8e799..ebcb9b843 100644
--- a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
+++ b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImage.java
@@ -39,7 +39,9 @@ import org.opencv.imgproc.Imgproc;
 import java.awt.Color;
 import java.awt.Graphics2D;
 import java.awt.image.BufferedImage;
+import java.awt.image.DataBuffer;
 import java.awt.image.DataBufferByte;
+import java.awt.image.DataBufferInt;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.ByteBuffer;
@@ -75,10 +77,41 @@ class OpenCVImage implements Image {
 
     /** {@inheritDoc} */
     @Override
-    public Object getWrappedImage() {
+    public Mat getWrappedImage() {
         return image;
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public OpenCVImage resize(int width, int height, boolean copy) {
+        if (!copy && image.width() == width && image.height() == height) {
+            return this;
+        }
+
+        Mat resized = new Mat();
+        Imgproc.resize(image, resized, new Size(width, height));
+        return new OpenCVImage(resized);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Image getMask(int[][] mask) {
+        int w = mask[0].length;
+        int h = mask.length;
+        OpenCVImage resized = resize(w, h, false);
+        Mat img = resized.getWrappedImage();
+        Mat ret = new Mat(h, w, CvType.CV_8UC4);
+        for (int y = 0; y < h; ++y) {
+            for (int x = 0; x < w; ++x) {
+                if (mask[y][x] != 0) {
+                    double[] data = img.get(y, x);
+                    ret.put(y, x, data[0], data[1], data[2], 255);
+                }
+            }
+        }
+        return new OpenCVImage(ret);
+    }
+
     /** {@inheritDoc} */
     @Override
     public Image getSubImage(int x, int y, int w, int h) {
@@ -185,6 +218,50 @@ class OpenCVImage implements Image {
         }
     }
 
+    /** {@inheritDoc} */
+    @Override
+    public void drawImage(Image overlay, boolean resize) {
+        if (!(overlay instanceof OpenCVImage)) {
+            throw new IllegalArgumentException("Only OpenCVImage allowed");
+        }
+        if (resize) {
+            overlay = overlay.resize(getWidth(), getHeight(), false);
+        }
+        Mat mat = (Mat) overlay.getWrappedImage();
+        if (mat.elemSize() != 4) {
+            mat.copyTo(image);
+            return;
+        }
+        int w = Math.min(image.width(), mat.width());
+        int h = Math.min(image.height(), mat.height());
+        for (int y = 0; y < h; ++y) {
+            for (int x = 0; x < w; ++x) {
+                /*
+                 * RA = SA + DA × (1 − SA)
+                 * R[0] = (S[0]×SA + D[0]×DA×(1 − SA)) / RA
+                 */
+                double[] src = mat.get(y, x);
+                double[] dest = image.get(y, x);
+                double sa = src[3];
+                double da;
+                double ra;
+                if (dest.length == 3) {
+                    da = 255 - sa;
+                    ra = 255;
+                } else {
+                    da = dest[3] * (255 - sa) / 255;
+                    ra = sa + da;
+                    dest[3] = ra;
+                }
+
+                dest[0] = (src[0] * sa + dest[0] * da) / ra;
+                dest[1] = (src[1] * sa + dest[1] * da) / ra;
+                dest[2] = (src[2] * sa + dest[2] * da) / ra;
+                image.put(y, x, dest);
+            }
+        }
+    }
+
     /** {@inheritDoc} */
     @Override
     public List<BoundingBox> findBoundingBoxes() {
@@ -260,8 +337,21 @@ class OpenCVImage implements Image {
     private static Mat image2Mat(BufferedImage img) {
         int width = img.getWidth();
         int height = img.getHeight();
-        byte[] data = ((DataBufferByte) img.getRaster().getDataBuffer()).getData();
-        Mat mat = new Mat(height, width, CvType.CV_8UC3);
+        byte[] data;
+        Mat mat;
+        DataBuffer buf = img.getRaster().getDataBuffer();
+        if (buf instanceof DataBufferByte) {
+            data = ((DataBufferByte) buf).getData();
+            mat = new Mat(height, width, CvType.CV_8UC3);
+        } else if (buf instanceof DataBufferInt) {
+            int[] intData = ((DataBufferInt) buf).getData();
+            data = new byte[intData.length * 4];
+            ByteBuffer bb = ByteBuffer.wrap(data);
+            bb.asIntBuffer().put(intData);
+            mat = new Mat(height, width, CvType.CV_8UC4);
+        } else {
+            throw new IllegalArgumentException("Unsupported image type: " + buf.getClass());
+        }
         mat.put(0, 0, data);
         return mat;
     }
diff --git a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java
index 2b5baa3e9..9aef084ca 100644
--- a/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java
+++ b/extensions/opencv/src/main/java/ai/djl/opencv/OpenCVImageFactory.java
@@ -30,6 +30,8 @@ import org.opencv.imgproc.Imgproc;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.nio.ByteBuffer;
+import java.nio.IntBuffer;
 import java.nio.file.Path;
 
 /** {@code OpenCVImageFactory} is a high performance implementation of {@link ImageFactory}. */
@@ -46,6 +48,7 @@ public class OpenCVImageFactory extends ImageFactory {
     /** {@inheritDoc} */
     @Override
     public Image fromFile(Path path) throws IOException {
+        // Load image without alpha channel
         Mat img = Imgcodecs.imread(path.toAbsolutePath().toString());
         if (img.empty()) {
             throw new IOException("Read image failed: " + path);
@@ -99,4 +102,20 @@ public class OpenCVImageFactory extends ImageFactory {
         Imgproc.cvtColor(img, img, Imgproc.COLOR_RGB2BGR);
         return new OpenCVImage(img);
     }
+
+    /** {@inheritDoc} */
+    @Override
+    public Image fromPixels(int[] pixels, int width, int height) {
+        Mat img = new Mat(height, width, CvType.CV_8UC4);
+        byte[] data = new byte[width * height * 4];
+        IntBuffer buf = ByteBuffer.wrap(data).asIntBuffer();
+        for (int pixel : pixels) {
+            int r = pixel >> 8 & 0xff00;
+            int g = pixel << 8 & 0xff0000;
+            int b = pixel << 24 & 0xff000000;
+            buf.put(pixel >>> 24 | b | g | r);
+        }
+        img.put(0, 0, data);
+        return new OpenCVImage(img);
+    }
 }
diff --git a/extensions/tablesaw/src/main/java/ai/djl/tablesaw/TablesawDataset.java b/extensions/tablesaw/src/main/java/ai/djl/tablesaw/TablesawDataset.java
index aaa277255..183775154 100644
--- a/extensions/tablesaw/src/main/java/ai/djl/tablesaw/TablesawDataset.java
+++ b/extensions/tablesaw/src/main/java/ai/djl/tablesaw/TablesawDataset.java
@@ -104,7 +104,6 @@ public class TablesawDataset extends TabularDataset {
          * @return the new {@link TablesawDataset}
          */
         public TablesawDataset build() {
-            validate();
             return new TablesawDataset(this);
         }
     }
diff --git a/extensions/tokenizers/src/main/python/arg_parser.py b/extensions/tokenizers/src/main/python/arg_parser.py
index eb88acb8e..8af923ba2 100644
--- a/extensions/tokenizers/src/main/python/arg_parser.py
+++ b/extensions/tokenizers/src/main/python/arg_parser.py
@@ -27,6 +27,10 @@ def converter_args():
                         "--retry-failed",
                         action='store_true',
                         help="Retry failed model")
+    parser.add_argument("-u",
+                        "--cpu-only",
+                        action='store_true',
+                        help="Only validate jit traced model on CPU")
     group = parser.add_mutually_exclusive_group(required=True)
     group.add_argument(
         "-c",
diff --git a/extensions/tokenizers/src/main/python/huggingface_converter.py b/extensions/tokenizers/src/main/python/huggingface_converter.py
index 1eb521b40..ec1aeebd4 100644
--- a/extensions/tokenizers/src/main/python/huggingface_converter.py
+++ b/extensions/tokenizers/src/main/python/huggingface_converter.py
@@ -14,6 +14,7 @@
 import logging
 import os.path
 import shutil
+from argparse import Namespace
 
 import torch
 from huggingface_hub import hf_hub_download
@@ -35,7 +36,7 @@ class HuggingfaceConverter:
         self.inputs = None
         self.outputs = None
 
-    def save_model(self, model_info, output_dir: str, temp_dir: str):
+    def save_model(self, model_info, args: Namespace, temp_dir: str):
         model_id = model_info.modelId
         if not os.path.exists(temp_dir):
             os.makedirs(temp_dir)
@@ -56,7 +57,7 @@ class HuggingfaceConverter:
             return False, "Failed to trace model", -1
 
         result, reason = self.verify_jit_model(hf_pipeline, model_file,
-                                               include_types)
+                                               include_types, args.cpu_only)
         if not result:
             include_types = True
             model_file = self.jit_trace_model(hf_pipeline, model_id, temp_dir,
@@ -65,11 +66,12 @@ class HuggingfaceConverter:
                 return False, reason, -1
 
             result, reason = self.verify_jit_model(hf_pipeline, model_file,
-                                                   include_types)
+                                                   include_types,
+                                                   args.cpu_only)
             if not result:
                 return False, reason, -1
 
-        size = self.save_to_model_zoo(model_info, output_dir, temp_dir,
+        size = self.save_to_model_zoo(model_info, args.output_dir, temp_dir,
                                       hf_pipeline, include_types)
 
         return True, None, size
@@ -157,7 +159,7 @@ class HuggingfaceConverter:
         return file_size
 
     def verify_jit_model(self, hf_pipeline, model_file: str,
-                         include_types: bool):
+                         include_types: bool, cpu_only: bool):
         logging.info(
             f"Verifying torchscript model(include_token_types={include_types}): {model_file} ..."
         )
@@ -168,7 +170,7 @@ class HuggingfaceConverter:
         input_ids = encoding["input_ids"]
         attention_mask = encoding["attention_mask"]
         token_type_ids = encoding.get("token_type_ids")
-        if torch.cuda.is_available():
+        if torch.cuda.is_available() and not cpu_only:
             traced_model = torch.jit.load(model_file, map_location='cuda:0')
             traced_model.to(self.device)
             input_ids = input_ids.to(self.device)
diff --git a/extensions/tokenizers/src/main/python/huggingface_models.py b/extensions/tokenizers/src/main/python/huggingface_models.py
index 6f974febb..3418815d5 100644
--- a/extensions/tokenizers/src/main/python/huggingface_models.py
+++ b/extensions/tokenizers/src/main/python/huggingface_models.py
@@ -124,7 +124,7 @@ class HuggingfaceModels:
         return ret
 
     def update_progress(self, model_info: ModelInfo, application: str,
-                        result: bool, reason: str, size: int):
+                        result: bool, reason: str, size: int, cpu_only: bool):
         status = {
             "result": "success" if result else "failed",
             "application": application,
@@ -134,6 +134,9 @@ class HuggingfaceModels:
         }
         if reason:
             status["reason"] = reason
+        if cpu_only:
+            status["cpu_only"] = True
+
         self.processed_models[model_info.modelId] = status
 
         dict_file = os.path.join(self.output_dir, "models.json")
diff --git a/extensions/tokenizers/src/main/python/model_zoo_importer.py b/extensions/tokenizers/src/main/python/model_zoo_importer.py
index a4e2c6d86..2841dbde4 100644
--- a/extensions/tokenizers/src/main/python/model_zoo_importer.py
+++ b/extensions/tokenizers/src/main/python/model_zoo_importer.py
@@ -47,13 +47,12 @@ def main():
         model_info = model["model_info"]
         converter = SUPPORTED_TASK[task]
 
-        result, reason, size = converter.save_model(model_info,
-                                                    args.output_dir, temp_dir)
+        result, reason, size = converter.save_model(model_info, args, temp_dir)
         if not result:
             logging.error(f"{model_info.modelId}: {reason}")
 
         huggingface_models.update_progress(model_info, converter.application,
-                                           result, reason, size)
+                                           result, reason, size, args.cpu_only)
         shutil.rmtree(temp_dir)
 
     logging.info("finished.")
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/object_detection/yolo/YOLOV3.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/object_detection/yolo/YOLOV3.java
new file mode 100644
index 000000000..1a329cf54
--- /dev/null
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/object_detection/yolo/YOLOV3.java
@@ -0,0 +1,533 @@
+/*
+ * 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.object_detection.yolo;
+
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.util.NDImageUtils;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDArrays;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.AbstractBlock;
+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.core.Linear;
+import ai.djl.nn.norm.BatchNorm;
+import ai.djl.nn.pooling.Pool;
+import ai.djl.training.ParameterStore;
+import ai.djl.util.PairList;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+
+/**
+ * {@code YOLOV3} contains a generic implementation of yolov3 adapted from
+ * https://github.com/bubbliiiing/yolo3-pytorch (Original author bubbliiiing). Yolov3 is a fast and
+ * accurate model for ObjectDetection tasks.
+ *
+ * @see <a href="https://arxiv.org/abs/1804.02767">for more information and knowledge about
+ *     yolov3</a>
+ */
+public final class YOLOV3 extends AbstractBlock {
+    private static final byte VERSION = 1;
+    private SequentialBlock darkNet53; // backBone of YOLOv3
+
+    private Block lastLayer0;
+    private Block layer0Output;
+
+    private Block lastLayer1Conv;
+    private Block lastLayer1UpSample;
+    private Block lastLayer1;
+    private Block layer1Output;
+
+    private Block lastLayer2Conv;
+    private Block lastLayer2UpSample;
+    private Block lastLayer2;
+    private Block layer2Output;
+
+    static final int[] REPEATS = {
+        1, 2, 8, 8, 4
+    }; // the repeat times of the darkNet53 residual units
+
+    static final int[] FILTERS = {
+        32, 64, 128, 256, 512, 1024
+    }; // the filters of darkNet53 residual units
+
+    private YOLOV3(Builder builder) {
+        super(VERSION);
+        darkNet53 = addChildBlock("darkNet53", darkNet53(builder, true));
+        lastLayer0 =
+                addChildBlock(
+                        "lastLayer0",
+                        makeLastLayers(
+                                FILTERS[4],
+                                FILTERS[5],
+                                builder.batchNormMomentum,
+                                builder.leakyAlpha));
+        layer0Output =
+                addChildBlock(
+                        "layer0Output",
+                        makeOutputLayers(
+                                FILTERS[5],
+                                3 * (builder.numClasses + 5),
+                                builder.batchNormMomentum,
+                                builder.leakyAlpha));
+
+        lastLayer1Conv =
+                addChildBlock(
+                        "lastLayer1Conv",
+                        convolutionBlock(256, 1, builder.batchNormMomentum, builder.leakyAlpha));
+        lastLayer1UpSample = addChildBlock("lastLayer1UpSample", upSampleBlockNearest());
+        lastLayer1 =
+                addChildBlock(
+                        "lastLayer1",
+                        makeLastLayers(
+                                FILTERS[3],
+                                FILTERS[4],
+                                builder.batchNormMomentum,
+                                builder.leakyAlpha));
+        layer1Output =
+                addChildBlock(
+                        "layer1Output",
+                        makeOutputLayers(
+                                FILTERS[4],
+                                3 * (builder.numClasses + 5),
+                                builder.batchNormMomentum,
+                                builder.leakyAlpha));
+
+        lastLayer2Conv =
+                addChildBlock(
+                        "lastLayer2Conv",
+                        convolutionBlock(128, 1, builder.batchNormMomentum, builder.leakyAlpha));
+        lastLayer2UpSample = addChildBlock("lastLayer2UpSample", upSampleBlockNearest());
+        lastLayer2 =
+                addChildBlock(
+                        "lastLayer2",
+                        makeLastLayers(
+                                FILTERS[2],
+                                FILTERS[3],
+                                builder.batchNormMomentum,
+                                builder.leakyAlpha));
+        layer2Output =
+                addChildBlock(
+                        "layer2Output",
+                        makeOutputLayers(
+                                FILTERS[3],
+                                3 * (builder.numClasses + 5),
+                                builder.batchNormMomentum,
+                                builder.leakyAlpha));
+    }
+
+    /**
+     * Builds a {@link Block} that represents an upSampleLayer(the nearest mode) for yolov3.
+     *
+     * @return a {@link Block} that represent an upSampleLayer for yolov3
+     */
+    public static Block upSampleBlockNearest() {
+        // transpose + upSample + transpose
+        return new SequentialBlock()
+                .addSingleton(array -> array.transpose(0, 2, 3, 1))
+                .addSingleton(
+                        array ->
+                                NDImageUtils.resize(
+                                        array,
+                                        (int) (array.getShape().get(1) * 2),
+                                        (int) (array.getShape().get(2) * 2),
+                                        Image.Interpolation.NEAREST))
+                .addSingleton(array -> array.transpose(0, 3, 1, 2));
+    }
+
+    /**
+     * Builds a {@link Block} that represents a conv-bn-leakyRelu unit for darkNet53.
+     *
+     * @param filters the number of filters for conv
+     * @param kernel the kernel size for conv
+     * @param batchNormMomentum the momentum for batchNorm layer
+     * @param leakyAlpha the alpha for leakyRelu activation
+     * @return a {@link Block} that represents a conv-bn-leakyRelu unit for darkNet53
+     */
+    public static Block convolutionBlock(
+            int filters, int kernel, float batchNormMomentum, float leakyAlpha) {
+        int pad = 0;
+        if (kernel > 0) {
+            pad = (kernel - 1) >> 1;
+        }
+        return new SequentialBlock()
+                .add(
+                        Conv2d.builder()
+                                .setFilters(filters)
+                                .setKernelShape(new Shape(kernel, kernel))
+                                .optPadding(new Shape(pad, pad))
+                                .build())
+                .add(BatchNorm.builder().optEpsilon(2E-5f).optMomentum(batchNormMomentum).build())
+                .add(Activation.leakyReluBlock(leakyAlpha));
+    }
+
+    /**
+     * Builds a {@link Block} that represents the feature head in yolov3.
+     *
+     * @param filtersIn the number of input filters
+     * @param filtersOut the number of output filters
+     * @param batchNormMomentum the momentum of batchNorm layer
+     * @param leakyAlpha the alpha value for leakyRelu activation
+     * @return a {@link Block} that represents the feature head in yolov3.
+     */
+    public static Block makeLastLayers(
+            int filtersIn, int filtersOut, float batchNormMomentum, float leakyAlpha) {
+        return new SequentialBlock()
+                .add(convolutionBlock(filtersIn, 1, batchNormMomentum, leakyAlpha))
+                .add(convolutionBlock(filtersOut, 3, batchNormMomentum, leakyAlpha))
+                .add(convolutionBlock(filtersIn, 1, batchNormMomentum, leakyAlpha))
+                .add(convolutionBlock(filtersOut, 3, batchNormMomentum, leakyAlpha))
+                .add(convolutionBlock(filtersIn, 1, batchNormMomentum, leakyAlpha));
+    }
+
+    /**
+     * Builds a {@link Block} that represents the output layer of yolov3.
+     *
+     * @param filtersOut the number of output filters
+     * @param outClass the number of output classes
+     * @param batchNormMomentum the momentum for batchNorm layer
+     * @param leakyAlpha the alpha for leakyRelu activation
+     * @return a {@link Block} that represents the output layer of yolov3.
+     */
+    public static Block makeOutputLayers(
+            int filtersOut, int outClass, float batchNormMomentum, float leakyAlpha) {
+        return new SequentialBlock()
+                .add(convolutionBlock(filtersOut, 3, batchNormMomentum, leakyAlpha))
+                .add(Conv2d.builder().setFilters(outClass).setKernelShape(new Shape(1, 1)).build());
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected NDList forwardInternal(
+            ParameterStore parameterStore,
+            NDList inputs,
+            boolean training,
+            PairList<String, Object> params) {
+        NDList xList = darkNet53.forward(parameterStore, inputs, training);
+        NDArray x0 = xList.get(7);
+        NDArray x1 = xList.get(6);
+        NDArray x2 = xList.get(5);
+
+        // the first feature layer
+        NDList out0Branch = lastLayer0.forward(parameterStore, new NDList(x0), training);
+        NDList out0 = layer0Output.forward(parameterStore, out0Branch, training);
+        NDList x1In = lastLayer1Conv.forward(parameterStore, out0Branch, training);
+        x1In = lastLayer1UpSample.forward(parameterStore, x1In, training);
+        x1In = new NDList(x1In.singletonOrThrow().concat(x1, 1));
+
+        NDList out1Branch = lastLayer1.forward(parameterStore, x1In, training);
+        NDList out1 = layer1Output.forward(parameterStore, out1Branch, training);
+
+        NDList x2In = lastLayer2Conv.forward(parameterStore, out1Branch, training);
+        x2In = lastLayer2UpSample.forward(parameterStore, x2In, training);
+        x2In = new NDList(x2In.singletonOrThrow().concat(x2, 1));
+
+        // the third feature layer
+        NDList out2 = lastLayer2.forward(parameterStore, x2In, training);
+        out2 = layer2Output.forward(parameterStore, out2, training);
+
+        // Outputs
+        return new NDList(
+                out0.singletonOrThrow(), out1.singletonOrThrow(), out2.singletonOrThrow());
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Shape[] getOutputShapes(Shape[] inputShapes) {
+        Shape[] current = inputShapes;
+        Shape[] outputs = new Shape[3];
+        Shape[] darkNetOutputs = new Shape[8];
+        int index = 0;
+        for (String name : children.keys()) {
+            Block block = children.get(name);
+            if (name.contains("darkNet")) {
+                darkNetOutputs = block.getOutputShapes(current);
+                current = new Shape[] {darkNetOutputs[7]};
+            } else if (name.contains("lastLayer")) {
+                if ("05lastLayer1UpSample".equals(name)) {
+                    current = block.getOutputShapes(current);
+                    current =
+                            new Shape[] {
+                                new Shape(
+                                        current[0].get(0),
+                                        current[0].get(1) + darkNetOutputs[6].get(1),
+                                        current[0].get(2),
+                                        current[0].get(3))
+                            };
+                } else if ("09lastLayer2UpSample".equals(name)) {
+                    current = block.getOutputShapes(current);
+                    current =
+                            new Shape[] {
+                                new Shape(
+                                        current[0].get(0),
+                                        current[0].get(1) + darkNetOutputs[5].get(1),
+                                        current[0].get(2),
+                                        current[0].get(3))
+                            };
+                } else {
+                    current = block.getOutputShapes(current);
+                }
+            } else if (!name.contains("Output")) {
+                current = block.getOutputShapes(current);
+            } else { // name.contains("Output")
+                Shape[] output = block.getOutputShapes(current);
+                outputs[index++] = output[0];
+            }
+        }
+        return outputs;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void initializeChildBlocks(NDManager manager, DataType dataType, Shape... inputShapes) {
+        Shape[] current = inputShapes;
+        Shape[] darkNetOutputs = new Shape[8];
+        for (String name : children.keys()) {
+            Block block = children.get(name);
+            block.initialize(manager, dataType, current);
+            if (name.contains("darkNet")) {
+                darkNetOutputs = block.getOutputShapes(current);
+                current = new Shape[] {darkNetOutputs[7]};
+            } else if (name.contains("lastLayer")) {
+                if ("05lastLayer1UpSample".equals(name)) {
+                    current = block.getOutputShapes(current);
+                    current =
+                            new Shape[] {
+                                new Shape(
+                                        current[0].get(0),
+                                        current[0].get(1) + darkNetOutputs[6].get(1),
+                                        current[0].get(2),
+                                        current[0].get(3))
+                            };
+                } else if ("09lastLayer2UpSample".equals(name)) {
+                    current = block.getOutputShapes(current);
+                    current =
+                            new Shape[] {
+                                new Shape(
+                                        current[0].get(0),
+                                        current[0].get(1) + darkNetOutputs[5].get(1),
+                                        current[0].get(2),
+                                        current[0].get(3))
+                            };
+                } else {
+                    current = block.getOutputShapes(current);
+                }
+            } else if (!name.contains("Output")) {
+                current = block.getOutputShapes(current);
+            } else { // name.contains("Output")
+                Shape[] output = block.getOutputShapes(current);
+            }
+        }
+    }
+
+    /**
+     * Builds a {@link Block} that a basic residual block unit used in DarkNet53.
+     *
+     * @param filters the output filter of the Convolutional Layer
+     * @param batchNormMomentum the momentum used for computing batchNorm
+     * @param leakyAlpha the alpha used in LeakyRelu Function
+     * @return a basic residual block unit used in DarkNet53
+     */
+    public static Block basicBlock(int filters, float batchNormMomentum, float leakyAlpha) {
+        SequentialBlock block = new SequentialBlock();
+        block.add(Conv2d.builder().setFilters(filters / 2).setKernelShape(new Shape(1, 1)).build())
+                .add(BatchNorm.builder().optEpsilon(2E-5f).optMomentum(batchNormMomentum).build())
+                .add(Activation.leakyReluBlock(leakyAlpha))
+                .add(
+                        Conv2d.builder()
+                                .setFilters(filters)
+                                .setKernelShape(new Shape(3, 3))
+                                .optPadding(new Shape(1, 1))
+                                .build())
+                .add(BatchNorm.builder().optEpsilon(2E-5f).optMomentum(batchNormMomentum).build())
+                .add(Activation.leakyReluBlock(leakyAlpha));
+        return new ParallelBlock(
+                list ->
+                        new NDList(
+                                NDArrays.add(
+                                        list.get(0).singletonOrThrow(),
+                                        list.get(1).singletonOrThrow())),
+                Arrays.asList(block, Blocks.identityBlock()));
+    }
+
+    /**
+     * Creates repeated Residual Blocks used in DarkNet53.
+     *
+     * @param filters the output filters of the final Convolutional Layer
+     * @param repeats the repeat times of a residual unit
+     * @param batchNormMomentum the momentum used for computing batchNorm
+     * @param leakyAlpha the alpha used in LeakyRelu Function
+     * @return several repeats of a residual block
+     */
+    public static Block makeLayer(
+            int filters, int repeats, float batchNormMomentum, float leakyAlpha) {
+        List<Block> layer = new ArrayList<>();
+        SequentialBlock convolutionalLayer = new SequentialBlock();
+        convolutionalLayer
+                .add(
+                        Conv2d.builder()
+                                .setFilters(filters)
+                                .setKernelShape(new Shape(3, 3))
+                                .optStride(new Shape(2, 2))
+                                .optPadding(new Shape(1, 1))
+                                .build())
+                .add(BatchNorm.builder().optEpsilon(2E-5f).optMomentum(batchNormMomentum).build())
+                .add(Activation.leakyReluBlock(leakyAlpha));
+        for (int i = 0; i < repeats; i++) {
+            layer.add(basicBlock(filters, batchNormMomentum, leakyAlpha));
+        }
+        return new SequentialBlock().add(convolutionalLayer).addAll(layer);
+    }
+
+    private static SequentialBlock darkNet53(Builder builder, boolean setReturnIntermediate) {
+        SequentialBlock darkNet53 = new SequentialBlock();
+        darkNet53.setReturnIntermediate(setReturnIntermediate); // return interMediate results;
+        darkNet53
+                .add(
+                        Conv2d.builder()
+                                .setFilters(FILTERS[0])
+                                .optPadding(new Shape(1, 1))
+                                .setKernelShape(new Shape(3, 3))
+                                .build())
+                .add(
+                        BatchNorm.builder()
+                                .optEpsilon(2E-5f)
+                                .optMomentum(builder.batchNormMomentum)
+                                .build())
+                .add(Activation.leakyReluBlock(builder.leakyAlpha))
+                .add(
+                        makeLayer(
+                                FILTERS[1],
+                                REPEATS[0],
+                                builder.batchNormMomentum,
+                                builder.leakyAlpha))
+                .add(
+                        makeLayer(
+                                FILTERS[2],
+                                REPEATS[1],
+                                builder.batchNormMomentum,
+                                builder.leakyAlpha))
+                .add(
+                        makeLayer(
+                                FILTERS[3],
+                                REPEATS[2],
+                                builder.batchNormMomentum,
+                                builder.leakyAlpha))
+                .add(
+                        makeLayer(
+                                FILTERS[4],
+                                REPEATS[3],
+                                builder.batchNormMomentum,
+                                builder.leakyAlpha))
+                .add(
+                        makeLayer(
+                                FILTERS[5],
+                                REPEATS[4],
+                                builder.batchNormMomentum,
+                                builder.leakyAlpha));
+        return darkNet53;
+    }
+
+    /**
+     * Creates a builder to build a {@link YOLOV3}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** The Builder to construct a {@link YOLOV3} object. */
+    public static final class Builder {
+        int numClasses = 20;
+        float batchNormMomentum = 0.9f;
+        float leakyAlpha = 0.1f;
+        int darkNetOutSize = 10;
+
+        /**
+         * Sets the number of classes for yolov3.
+         *
+         * @param numClasses the number of classes
+         * @return @return this {@code Builder}
+         */
+        public Builder setNumClasses(int numClasses) {
+            this.numClasses = numClasses;
+            return this;
+        }
+
+        /**
+         * Sets the momentum for batchNorm layer.
+         *
+         * @param batchNormMomentum the momentum for batchNorm layer
+         * @return this {@code Builder}
+         */
+        public Builder optBatchNormMomentum(float batchNormMomentum) {
+            this.batchNormMomentum = batchNormMomentum;
+            return this;
+        }
+
+        /**
+         * Sets the alpha for leakyRelu activation.
+         *
+         * @param leakyAlpha the alpha for leakyRelu activation
+         * @return this {@code Builder}
+         */
+        public Builder optLeakyAlpha(float leakyAlpha) {
+            this.leakyAlpha = leakyAlpha;
+            return this;
+        }
+
+        /**
+         * Sets the out size of darkNet for testing.
+         *
+         * @param darkNetOutSize the out size of darkNet
+         * @return this {@code Builder}
+         */
+        public Builder optDarkNetOutSize(int darkNetOutSize) {
+            this.darkNetOutSize = darkNetOutSize;
+            return this;
+        }
+
+        /**
+         * Builds a {@link YOLOV3} block.
+         *
+         * @return a {@link YOLOV3} block
+         */
+        public Block build() {
+            return new YOLOV3(this);
+        }
+
+        /**
+         * Builds a {@link Block} that represents the backbone of yolov3, which is called DarkNet53.
+         * This can be used for testing and transfer learning.
+         *
+         * @return a {@link Block} that represents darkNet53
+         */
+        public Block buildDarkNet() {
+            Block block = darkNet53(this, false);
+            return new SequentialBlock()
+                    .add(block)
+                    .add(Pool.globalAvgPool2dBlock())
+                    .add(Linear.builder().setUnits(darkNetOutSize).build());
+        }
+    }
+}
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/object_detection/yolo/package-info.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/object_detection/yolo/package-info.java
new file mode 100644
index 000000000..5b25239de
--- /dev/null
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/cv/object_detection/yolo/package-info.java
@@ -0,0 +1,18 @@
+/*
+ * 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.
+ */
+
+/**
+ * Contains classes for the {@link ai.djl.Application.CV#OBJECT_DETECTION} models in the {@link
+ * ai.djl.basicmodelzoo.BasicModelZoo}.
+ */
+package ai.djl.basicmodelzoo.cv.object_detection.yolo;
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/TabNet.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/TabNet.java
index dac0827ff..be2f8d007 100644
--- a/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/TabNet.java
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/TabNet.java
@@ -69,10 +69,7 @@ public final class TabNet extends AbstractBlock {
         batchNorm =
                 addChildBlock(
                         "batchNorm",
-                        BatchNorm.builder()
-                                .optAxis(0)
-                                .optMomentum(builder.batchNormMomentum)
-                                .build());
+                        BatchNorm.builder().optMomentum(builder.batchNormMomentum).build());
         List<Block> sharedBlocks = new ArrayList<>();
         for (int i = 0; i < builder.numShared; i++) {
             sharedBlocks.add(
@@ -169,7 +166,6 @@ public final class TabNet extends AbstractBlock {
         featureBlock
                 .add(
                         GhostBatchNorm.builder()
-                                .optAxis(0)
                                 .optVirtualBatchSize(virtualBatchSize)
                                 .optMomentum(batchNormMomentum)
                                 .build())
@@ -237,6 +233,7 @@ public final class TabNet extends AbstractBlock {
             PairList<String, Object> params) {
         NDManager manager = inputs.getManager();
         NDArray input = inputs.singletonOrThrow();
+        input = input.reshape(input.size(0), input.size() / input.size(0)); // batch flatten
         NDArray x =
                 batchNorm.forward(parameterStore, new NDList(input), training).singletonOrThrow();
         NDArray xa =
@@ -347,11 +344,10 @@ public final class TabNet extends AbstractBlock {
                     addChildBlock(
                             "ghostBatchNorm",
                             GhostBatchNorm.builder()
-                                    .optAxis(0)
                                     .optVirtualBatchSize(virtualBatchSize)
                                     .optMomentum(batchNormMomentum)
                                     .build());
-            sparseMax = addChildBlock("sparseMax", new SparseMax(-1, 8));
+            sparseMax = addChildBlock("sparseMax", new SparseMax());
         }
 
         /** {@inheritDoc} */
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/package-info.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/package-info.java
index ff3201da0..3046799c8 100644
--- a/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/package-info.java
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/package-info.java
@@ -12,7 +12,7 @@
  */
 
 /**
- * Contains classes for the computer vision section ({@link ai.djl.Application.Tabular}) of the
- * {@link ai.djl.basicmodelzoo.BasicModelZoo}.
+ * Contains classes for the tabular section ({@link ai.djl.Application.Tabular}) of the {@link
+ * ai.djl.basicmodelzoo.BasicModelZoo}.
  */
 package ai.djl.basicmodelzoo.tabular;


Test patch:
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/BigGANTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/BigGANTranslatorFactoryTest.java
new file mode 100644
index 000000000..172354a88
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/translator/BigGANTranslatorFactoryTest.java
@@ -0,0 +1,55 @@
+/*
+ * 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.modality.cv.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.translate.Translator;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class BigGANTranslatorFactoryTest {
+
+    private BigGANTranslatorFactory factory;
+
+    @BeforeClass
+    public void setUp() {
+        factory = new BigGANTranslatorFactory();
+    }
+
+    @Test
+    public void testGetSupportedTypes() {
+        Assert.assertEquals(factory.getSupportedTypes().size(), 1);
+    }
+
+    @Test
+    public void testNewInstance() {
+        Map<String, Float> arguments = new HashMap<>();
+        arguments.put("truncation", 0.6f);
+        try (Model model = Model.newInstance("test")) {
+            Translator<int[], Image[]> translator =
+                    factory.newInstance(int[].class, Image[].class, model, arguments);
+            Assert.assertTrue(translator instanceof BigGANTranslator);
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(Input.class, Output.class, model, arguments));
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactoryTest.java
new file mode 100644
index 000000000..941e2c3dc
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/translator/ImageClassificationTranslatorFactoryTest.java
@@ -0,0 +1,78 @@
+/*
+ * 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.modality.cv.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.Classifications;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.translator.wrapper.FileTranslator;
+import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
+import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.translate.Translator;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.io.InputStream;
+import java.net.URL;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.Map;
+
+public class ImageClassificationTranslatorFactoryTest {
+
+    private ImageClassificationTranslatorFactory factory;
+
+    @BeforeClass
+    public void setUp() {
+        factory = new ImageClassificationTranslatorFactory();
+    }
+
+    @Test
+    public void testGetSupportedTypes() {
+        Assert.assertEquals(factory.getSupportedTypes().size(), 5);
+    }
+
+    @Test
+    public void testNewInstance() {
+        Map<String, String> arguments = new HashMap<>();
+        try (Model model = Model.newInstance("test")) {
+            Translator<Image, Classifications> translator1 =
+                    factory.newInstance(Image.class, Classifications.class, model, arguments);
+            Assert.assertTrue(translator1 instanceof ImageClassificationTranslator);
+
+            Translator<Path, Classifications> translator2 =
+                    factory.newInstance(Path.class, Classifications.class, model, arguments);
+            Assert.assertTrue(translator2 instanceof FileTranslator);
+
+            Translator<URL, Classifications> translator3 =
+                    factory.newInstance(URL.class, Classifications.class, model, arguments);
+            Assert.assertTrue(translator3 instanceof UrlTranslator);
+
+            Translator<InputStream, Classifications> translator4 =
+                    factory.newInstance(InputStream.class, Classifications.class, model, arguments);
+            Assert.assertTrue(translator4 instanceof InputStreamTranslator);
+
+            Translator<Input, Output> translator5 =
+                    factory.newInstance(Input.class, Output.class, model, arguments);
+            Assert.assertTrue(translator5 instanceof ImageServingTranslator);
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(Image.class, Output.class, model, arguments));
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactoryTest.java
new file mode 100644
index 000000000..79df6bcc9
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/translator/InstanceSegmentationTranslatorFactoryTest.java
@@ -0,0 +1,78 @@
+/*
+ * 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.modality.cv.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.translator.wrapper.FileTranslator;
+import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
+import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.translate.Translator;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.io.InputStream;
+import java.net.URL;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.Map;
+
+public class InstanceSegmentationTranslatorFactoryTest {
+
+    private InstanceSegmentationTranslatorFactory factory;
+
+    @BeforeClass
+    public void setUp() {
+        factory = new InstanceSegmentationTranslatorFactory();
+    }
+
+    @Test
+    public void testGetSupportedTypes() {
+        Assert.assertEquals(factory.getSupportedTypes().size(), 5);
+    }
+
+    @Test
+    public void testNewInstance() {
+        Map<String, String> arguments = new HashMap<>();
+        try (Model model = Model.newInstance("test")) {
+            Translator<Image, DetectedObjects> translator1 =
+                    factory.newInstance(Image.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator1 instanceof InstanceSegmentationTranslator);
+
+            Translator<Path, DetectedObjects> translator2 =
+                    factory.newInstance(Path.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator2 instanceof FileTranslator);
+
+            Translator<URL, DetectedObjects> translator3 =
+                    factory.newInstance(URL.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator3 instanceof UrlTranslator);
+
+            Translator<InputStream, DetectedObjects> translator4 =
+                    factory.newInstance(InputStream.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator4 instanceof InputStreamTranslator);
+
+            Translator<Input, Output> translator5 =
+                    factory.newInstance(Input.class, Output.class, model, arguments);
+            Assert.assertTrue(translator5 instanceof ImageServingTranslator);
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(Image.class, Output.class, model, arguments));
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactoryTest.java
new file mode 100644
index 000000000..ec54e3bcb
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/translator/SemanticSegmentationTranslatorFactoryTest.java
@@ -0,0 +1,55 @@
+/*
+ * 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.modality.cv.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.CategoryMask;
+import ai.djl.translate.Translator;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class SemanticSegmentationTranslatorFactoryTest {
+
+    private SemanticSegmentationTranslatorFactory factory;
+
+    @BeforeClass
+    public void setUp() {
+        factory = new SemanticSegmentationTranslatorFactory();
+    }
+
+    @Test
+    public void testGetSupportedTypes() {
+        Assert.assertEquals(factory.getSupportedTypes().size(), 1);
+    }
+
+    @Test
+    public void testNewInstance() {
+        Map<String, String> arguments = new HashMap<>();
+        try (Model model = Model.newInstance("test")) {
+            Translator<Image, CategoryMask> translator =
+                    factory.newInstance(Image.class, CategoryMask.class, model, arguments);
+            Assert.assertTrue(translator instanceof SemanticSegmentationTranslator);
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(Input.class, Output.class, model, arguments));
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactoryTest.java
new file mode 100644
index 000000000..0a2b54fd8
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/translator/SimplePoseTranslatorFactoryTest.java
@@ -0,0 +1,60 @@
+/*
+ * 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.modality.cv.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.Joints;
+import ai.djl.translate.Translator;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class SimplePoseTranslatorFactoryTest {
+
+    private SimplePoseTranslatorFactory factory;
+
+    @BeforeClass
+    public void setUp() {
+        factory = new SimplePoseTranslatorFactory();
+    }
+
+    @Test
+    public void testGetSupportedTypes() {
+        Assert.assertEquals(factory.getSupportedTypes().size(), 2);
+    }
+
+    @Test
+    public void testNewInstance() {
+        Map<String, String> arguments = new HashMap<>();
+        try (Model model = Model.newInstance("test")) {
+            Translator<Image, Joints> translator1 =
+                    factory.newInstance(Image.class, Joints.class, model, arguments);
+            Assert.assertTrue(translator1 instanceof SimplePoseTranslator);
+
+            Translator<Input, Output> translator2 =
+                    factory.newInstance(Input.class, Output.class, model, arguments);
+            Assert.assertTrue(translator2 instanceof ImageServingTranslator);
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(Image.class, Output.class, model, arguments));
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactoryTest.java
new file mode 100644
index 000000000..938a5cd0f
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/translator/SingleShotDetectionTranslatorFactoryTest.java
@@ -0,0 +1,73 @@
+/*
+ * 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.modality.cv.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.translator.wrapper.FileTranslator;
+import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
+import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.translate.Translator;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.io.InputStream;
+import java.net.URL;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.Map;
+
+public class SingleShotDetectionTranslatorFactoryTest {
+
+    private SingleShotDetectionTranslatorFactory factory;
+
+    @BeforeClass
+    public void setUp() {
+        factory = new SingleShotDetectionTranslatorFactory();
+    }
+
+    @Test
+    public void testNewInstance() {
+        Map<String, String> arguments = new HashMap<>();
+        try (Model model = Model.newInstance("test")) {
+            Translator<Image, DetectedObjects> translator1 =
+                    factory.newInstance(Image.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator1 instanceof SingleShotDetectionTranslator);
+
+            Translator<Path, DetectedObjects> translator2 =
+                    factory.newInstance(Path.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator2 instanceof FileTranslator);
+
+            Translator<URL, DetectedObjects> translator3 =
+                    factory.newInstance(URL.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator3 instanceof UrlTranslator);
+
+            Translator<InputStream, DetectedObjects> translator4 =
+                    factory.newInstance(InputStream.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator4 instanceof InputStreamTranslator);
+
+            Translator<Input, Output> translator5 =
+                    factory.newInstance(Input.class, Output.class, model, arguments);
+            Assert.assertTrue(translator5 instanceof ImageServingTranslator);
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(Image.class, Output.class, model, arguments));
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactoryTest.java
new file mode 100644
index 000000000..785ffa872
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/translator/StyleTransferTranslatorFactoryTest.java
@@ -0,0 +1,54 @@
+/*
+ * 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.modality.cv.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.translate.Translator;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.util.HashMap;
+import java.util.Map;
+
+public class StyleTransferTranslatorFactoryTest {
+
+    private StyleTransferTranslatorFactory factory;
+
+    @BeforeClass
+    public void setUp() {
+        factory = new StyleTransferTranslatorFactory();
+    }
+
+    @Test
+    public void testGetSupportedTypes() {
+        Assert.assertEquals(factory.getSupportedTypes().size(), 1);
+    }
+
+    @Test
+    public void testNewInstance() {
+        Map<String, String> arguments = new HashMap<>();
+        try (Model model = Model.newInstance("test")) {
+            Translator<Image, Image> translator =
+                    factory.newInstance(Image.class, Image.class, model, arguments);
+            Assert.assertTrue(translator instanceof StyleTransferTranslator);
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(Input.class, Output.class, model, arguments));
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/YoloTranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/YoloTranslatorFactoryTest.java
new file mode 100644
index 000000000..fae1a595f
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/translator/YoloTranslatorFactoryTest.java
@@ -0,0 +1,78 @@
+/*
+ * 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.modality.cv.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.translator.wrapper.FileTranslator;
+import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
+import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.translate.Translator;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.io.InputStream;
+import java.net.URL;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.Map;
+
+public class YoloTranslatorFactoryTest {
+
+    private YoloTranslatorFactory factory;
+
+    @BeforeClass
+    public void setUp() {
+        factory = new YoloTranslatorFactory();
+    }
+
+    @Test
+    public void testGetSupportedTypes() {
+        Assert.assertEquals(factory.getSupportedTypes().size(), 5);
+    }
+
+    @Test
+    public void testNewInstance() {
+        Map<String, String> arguments = new HashMap<>();
+        try (Model model = Model.newInstance("test")) {
+            Translator<Image, DetectedObjects> translator1 =
+                    factory.newInstance(Image.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator1 instanceof YoloTranslator);
+
+            Translator<Path, DetectedObjects> translator2 =
+                    factory.newInstance(Path.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator2 instanceof FileTranslator);
+
+            Translator<URL, DetectedObjects> translator3 =
+                    factory.newInstance(URL.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator3 instanceof UrlTranslator);
+
+            Translator<InputStream, DetectedObjects> translator4 =
+                    factory.newInstance(InputStream.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator4 instanceof InputStreamTranslator);
+
+            Translator<Input, Output> translator5 =
+                    factory.newInstance(Input.class, Output.class, model, arguments);
+            Assert.assertTrue(translator5 instanceof ImageServingTranslator);
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(Image.class, Output.class, model, arguments));
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactoryTest.java b/api/src/test/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactoryTest.java
new file mode 100644
index 000000000..a86c03fac
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/translator/YoloV5TranslatorFactoryTest.java
@@ -0,0 +1,78 @@
+/*
+ * 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.modality.cv.translator;
+
+import ai.djl.Model;
+import ai.djl.modality.Input;
+import ai.djl.modality.Output;
+import ai.djl.modality.cv.Image;
+import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.modality.cv.translator.wrapper.FileTranslator;
+import ai.djl.modality.cv.translator.wrapper.InputStreamTranslator;
+import ai.djl.modality.cv.translator.wrapper.UrlTranslator;
+import ai.djl.translate.Translator;
+
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import java.io.InputStream;
+import java.net.URL;
+import java.nio.file.Path;
+import java.util.HashMap;
+import java.util.Map;
+
+public class YoloV5TranslatorFactoryTest {
+
+    private YoloV5TranslatorFactory factory;
+
+    @BeforeClass
+    public void setUp() {
+        factory = new YoloV5TranslatorFactory();
+    }
+
+    @Test
+    public void testGetSupportedTypes() {
+        Assert.assertEquals(factory.getSupportedTypes().size(), 5);
+    }
+
+    @Test
+    public void testNewInstance() {
+        Map<String, String> arguments = new HashMap<>();
+        try (Model model = Model.newInstance("test")) {
+            Translator<Image, DetectedObjects> translator1 =
+                    factory.newInstance(Image.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator1 instanceof YoloV5Translator);
+
+            Translator<Path, DetectedObjects> translator2 =
+                    factory.newInstance(Path.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator2 instanceof FileTranslator);
+
+            Translator<URL, DetectedObjects> translator3 =
+                    factory.newInstance(URL.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator3 instanceof UrlTranslator);
+
+            Translator<InputStream, DetectedObjects> translator4 =
+                    factory.newInstance(InputStream.class, DetectedObjects.class, model, arguments);
+            Assert.assertTrue(translator4 instanceof InputStreamTranslator);
+
+            Translator<Input, Output> translator5 =
+                    factory.newInstance(Input.class, Output.class, model, arguments);
+            Assert.assertTrue(translator5 instanceof ImageServingTranslator);
+
+            Assert.assertThrows(
+                    IllegalArgumentException.class,
+                    () -> factory.newInstance(Image.class, Output.class, model, arguments));
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/modality/cv/translator/package-info.java b/api/src/test/java/ai/djl/modality/cv/translator/package-info.java
new file mode 100644
index 000000000..e94737154
--- /dev/null
+++ b/api/src/test/java/ai/djl/modality/cv/translator/package-info.java
@@ -0,0 +1,14 @@
+/*
+ * 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.
+ */
+/** Contains tests for {@link ai.djl.modality.cv.translator}. */
+package ai.djl.modality.cv.translator;
diff --git a/basicdataset/src/test/java/ai/djl/basicdataset/AmesRandomAccessTest.java b/basicdataset/src/test/java/ai/djl/basicdataset/AmesRandomAccessTest.java
index d0a008aeb..ab58b8cb8 100644
--- a/basicdataset/src/test/java/ai/djl/basicdataset/AmesRandomAccessTest.java
+++ b/basicdataset/src/test/java/ai/djl/basicdataset/AmesRandomAccessTest.java
@@ -64,6 +64,7 @@ public class AmesRandomAccessTest {
                             .addFeature("MiscVal")
                             .addFeature("id")
                             .addFeature("alley", true)
+                            .addNumericLabel("saleprice")
                             .setSampling(32, true)
                             .build();
 
@@ -75,7 +76,7 @@ public class AmesRandomAccessTest {
             Record record = amesRandomAccess.get(manager, 0);
             NDList data = record.getData();
             NDList labels = record.getLabels();
-            Assert.assertEquals(data.head().toFloatArray(), new float[] {0, 1, 1, 0, 0});
+            Assert.assertEquals(data.head().toFloatArray(), new float[] {0, 1, 0, 1, 0});
             Assert.assertEquals(labels.singletonOrThrow().toFloatArray()[0], 208500.0);
 
             try (Trainer trainer = model.newTrainer(config)) {
@@ -86,4 +87,39 @@ public class AmesRandomAccessTest {
             }
         }
     }
+
+    @Test
+    public void testAmesRandomAccessAllFeatures() throws IOException, TranslateException {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            AmesRandomAccess amesRandomAccess =
+                    AmesRandomAccess.builder()
+                            .optUsage(Dataset.Usage.TRAIN)
+                            .addAllFeatures()
+                            .setSampling(32, true)
+                            .build();
+
+            amesRandomAccess.prepare();
+
+            long size = amesRandomAccess.size();
+            Assert.assertEquals(size, 1460);
+
+            Record record = amesRandomAccess.get(manager, 0);
+            NDList data = record.getData();
+            Assert.assertEquals(data.head().size(), 79);
+        }
+    }
+
+    @Test
+    public void testAmesNoFeatures() {
+        Assert.assertThrows(
+                () -> {
+                    AmesRandomAccess amesRandomAccess =
+                            AmesRandomAccess.builder()
+                                    .optUsage(Dataset.Usage.TRAIN)
+                                    .setSampling(32, true)
+                                    .build();
+
+                    amesRandomAccess.prepare();
+                });
+    }
 }
diff --git a/basicdataset/src/test/java/ai/djl/basicdataset/DailyDelhiClimateTest.java b/basicdataset/src/test/java/ai/djl/basicdataset/DailyDelhiClimateTest.java
index a7efa6496..6a68c6431 100644
--- a/basicdataset/src/test/java/ai/djl/basicdataset/DailyDelhiClimateTest.java
+++ b/basicdataset/src/test/java/ai/djl/basicdataset/DailyDelhiClimateTest.java
@@ -51,6 +51,7 @@ public class DailyDelhiClimateTest {
                             .optUsage(Dataset.Usage.TEST)
                             .addFeature("humidity")
                             .addFeature("wind_speed")
+                            .noLabels()
                             .setSampling(32, true)
                             .build();
 
diff --git a/engines/pytorch/pytorch-model-zoo/src/test/resources/mlrepo/model/cv/semantic_segmentation/ai/djl/pytorch/deeplabv3/metadata.json b/engines/pytorch/pytorch-model-zoo/src/test/resources/mlrepo/model/cv/semantic_segmentation/ai/djl/pytorch/deeplabv3/metadata.json
index 31ddf1b81..0aa710262 100644
--- a/engines/pytorch/pytorch-model-zoo/src/test/resources/mlrepo/model/cv/semantic_segmentation/ai/djl/pytorch/deeplabv3/metadata.json
+++ b/engines/pytorch/pytorch-model-zoo/src/test/resources/mlrepo/model/cv/semantic_segmentation/ai/djl/pytorch/deeplabv3/metadata.json
@@ -31,9 +31,9 @@
       "files": {
         "model": {
           "uri": "0.0.1/deeplabv3.zip",
-          "sha1Hash": "06b8471f40d41d5f4fce7736cdd24081b9fa2e8e",
+          "sha1Hash": "417241c13140e8ad307fdb5884442183fb94e130",
           "name": "",
-          "size": 155844147
+          "size": 155844400
         }
       }
     }
diff --git a/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java b/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java
index 649ca193b..3e8937de7 100644
--- a/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java
+++ b/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java
@@ -19,11 +19,16 @@ import ai.djl.modality.Classifications;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
 import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.testing.TestRequirements;
 import ai.djl.training.util.ProgressBar;
+import ai.djl.translate.NoopTranslator;
 import ai.djl.translate.TranslateException;
 import ai.djl.util.Pair;
 
@@ -59,14 +64,21 @@ public class TfSsdTest {
                 Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
             Assert.assertEquals(model.describeInput().get(0).getValue(), new Shape(-1, -1, -1, 3));
             for (Pair<String, Shape> pair : model.describeOutput()) {
-                if (pair.getKey().contains("label")) {
-                    Assert.assertEquals(pair.getValue(), new Shape(-1, 1));
-                } else if (pair.getKey().contains("box")) {
-                    Assert.assertEquals(pair.getValue(), new Shape(-1, 4));
-                } else if (pair.getKey().contains("score")) {
-                    Assert.assertEquals(pair.getValue(), new Shape(-1, 1));
-                } else {
-                    throw new IllegalStateException("Unexpected output name:" + pair.getKey());
+                switch (pair.getKey()) {
+                    case "box":
+                    case "detection_boxes":
+                        Assert.assertEquals(pair.getValue(), new Shape(-1, 4));
+                        break;
+                    case "label":
+                    case "score":
+                    case "detection_class_entities":
+                    case "detection_class_labels":
+                    case "detection_class_names":
+                    case "detection_scores":
+                        Assert.assertEquals(pair.getValue(), new Shape(-1, 1));
+                        break;
+                    default:
+                        throw new IllegalStateException("Unexpected output name: " + pair.getKey());
                 }
             }
 
@@ -82,6 +94,32 @@ public class TfSsdTest {
         }
     }
 
+    @Test
+    public void testStringInputOutput() throws IOException, ModelException, TranslateException {
+        TestRequirements.notArm();
+
+        Criteria<NDList, NDList> criteria =
+                Criteria.builder()
+                        .optApplication(Application.CV.OBJECT_DETECTION)
+                        .setTypes(NDList.class, NDList.class)
+                        .optArtifactId("ssd")
+                        .optFilter("backbone", "mobilenet_v2")
+                        .optEngine("TensorFlow")
+                        .optProgress(new ProgressBar())
+                        .optTranslator(new NoopTranslator())
+                        .build();
+
+        try (ZooModel<NDList, NDList> model = criteria.loadModel();
+                Predictor<NDList, NDList> predictor = model.newPredictor();
+                NDManager manager = NDManager.newBaseManager()) {
+            NDArray array = manager.zeros(new Shape(1, 224, 224, 3));
+            NDList output = predictor.predict(new NDList(array));
+            Assert.assertEquals(output.size(), 5);
+            NDArray entities = output.get("detection_class_entities");
+            Assert.assertEquals(entities.getDataType(), DataType.STRING);
+        }
+    }
+
     private static void saveBoundingBoxImage(Image img, DetectedObjects detection)
             throws IOException {
         Path outputDir = Paths.get("build/output");
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java b/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
new file mode 100644
index 000000000..89a2dadcc
--- /dev/null
+++ b/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
@@ -0,0 +1,34 @@
+/*
+ * 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.examples.training;
+
+import ai.djl.testing.TestRequirements;
+import ai.djl.training.TrainingResult;
+import ai.djl.translate.TranslateException;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+
+public class TrainAirfoilWithTabNetTest {
+    @Test
+    public void testTrainAirfoilWithTabNet() throws TranslateException, IOException {
+        TestRequirements.engine("MXNet", "PyTorch");
+        String[] args = new String[] {"-g", "1", "-e", "20", "-b", "32"};
+        TrainingResult result = TrainAirfoilWithTabNet.runExample(args);
+        Assert.assertNotNull(result);
+        float loss = result.getValidateLoss();
+        Assert.assertTrue(loss < 50f, "Loss: " + loss);
+    }
+}
diff --git a/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java b/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java
index 917fa088e..69c683bcf 100644
--- a/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java
+++ b/extensions/opencv/src/test/java/ai/djl/opencv/OpenCVImageFactoryTest.java
@@ -25,9 +25,11 @@ import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDManager;
 import ai.djl.testing.TestRequirements;
 
+import org.opencv.core.Mat;
 import org.testng.Assert;
 import org.testng.annotations.Test;
 
+import java.awt.Color;
 import java.io.IOException;
 import java.io.OutputStream;
 import java.nio.file.Files;
@@ -156,4 +158,43 @@ public class OpenCVImageFactoryTest {
             }
         }
     }
+
+    @Test
+    public void testDrawImage() throws IOException {
+        ImageFactory factory = ImageFactory.getInstance();
+        int[] pixels = new int[64];
+        int index = 0;
+        for (int i = 0; i < 16; ++i) {
+            pixels[index++] = Color.RED.getRGB() & 0x7fffffff;
+        }
+        for (int i = 0; i < 16; ++i) {
+            pixels[index++] = Color.GREEN.getRGB() & 0x7fffffff;
+        }
+        for (int i = 0; i < 16; ++i) {
+            pixels[index++] = Color.BLUE.getRGB() & 0x7fffffff;
+        }
+        for (int i = 0; i < 16; ++i) {
+            pixels[index++] = Color.BLACK.getRGB() & 0x7fffffff;
+        }
+        Image img1 = factory.fromPixels(pixels, 8, 8);
+
+        pixels = new int[16];
+        for (int i = 0; i < 16; ++i) {
+            pixels[i] = Color.RED.getRGB() & 0x7fffffff;
+        }
+        Image img2 = factory.fromPixels(pixels, 2, 8);
+        Image mask = img2.getMask(new int[][] {{1, 1}, {1, 1}, {0, 0}, {0, 0}});
+        Image dup = img1.duplicate();
+        img1.drawImage(img2, false);
+        Mat mat = (Mat) img1.getWrappedImage();
+        byte[] data = new byte[64 * 4];
+        mat.get(0, 0, data);
+        Assert.assertEquals(data[2], -1); // red
+        Assert.assertEquals(data[3], -65); // alpha
+
+        dup.drawImage(mask, true);
+        ((Mat) dup.getWrappedImage()).get(0, 0, data);
+        Assert.assertEquals(data[2], -1); // red
+        Assert.assertEquals(data[3], -1); // alpha
+    }
 }
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/object_detection/YOLOv3Test.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/object_detection/YOLOv3Test.java
new file mode 100644
index 000000000..3eef4f6c6
--- /dev/null
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/object_detection/YOLOv3Test.java
@@ -0,0 +1,153 @@
+/*
+ * 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.object_detection;
+
+import ai.djl.Model;
+import ai.djl.basicmodelzoo.cv.object_detection.yolo.YOLOV3;
+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.testing.TestRequirements;
+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.training.loss.YOLOv3Loss;
+import ai.djl.translate.Batchifier;
+import ai.djl.util.PairList;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class YOLOv3Test {
+    @Test
+    public void testDarkNet() {
+        TrainingConfig config =
+                new DefaultTrainingConfig(Loss.softmaxCrossEntropyLoss())
+                        .optDevices(Engine.getInstance().getDevices(2))
+                        .optInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+        Block darkNet = YOLOV3.builder().buildDarkNet();
+        try (Model model = Model.newInstance("darkNet")) {
+            model.setBlock(darkNet);
+            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 = darkNet.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 expectedAtIndex6 =
+                        manager.ones(new Shape(64, 32, 3, 3)); // 64*32*3*3 for second layer
+                NDArray expectedAtIndex120 =
+                        manager.ones(new Shape(128, 256, 1, 1)); // 128*256*1*1  for 20th layer
+                NDArray expectedLinear =
+                        manager.ones(
+                                new Shape(
+                                        10,
+                                        1024)); // outSize = 10, linear layer size 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(6).getValue().getArray(), expectedAtIndex6);
+                Assertions.assertAlmostEquals(
+                        parameters.get(120).getValue().getArray(), expectedAtIndex120);
+                Assertions.assertAlmostEquals(
+                        parameters.get(312).getValue().getArray(), expectedLinear);
+            }
+        }
+    }
+
+    @Test
+    public void testYoloV3() {
+        TestRequirements.engine("PyTorch");
+
+        float[] anchorsArray = YOLOv3Loss.getPresetAnchors();
+        for (int i = 0; i < anchorsArray.length; i++) {
+            anchorsArray[i] = anchorsArray[i] * 256 / 416; // reshaping into the
+        }
+        TrainingConfig config =
+                new DefaultTrainingConfig(
+                                YOLOv3Loss.builder()
+                                        .setNumClasses(1)
+                                        .setInputShape(new Shape(256, 256))
+                                        .setAnchorsArray(anchorsArray)
+                                        .build())
+                        .optDevices(Engine.getInstance().getDevices(2))
+                        .optInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+
+        Block yolov3 = YOLOV3.builder().setNumClasses(1).build();
+        try (Model model = Model.newInstance("yolov3")) {
+            model.setBlock(yolov3);
+            try (Trainer trainer = model.newTrainer(config)) {
+                int batchSize = 4;
+                Shape inputShape = new Shape(batchSize, 3, 256, 256);
+                trainer.initialize(inputShape);
+                NDManager manager = trainer.getManager();
+                NDArray input = manager.ones(inputShape);
+                NDArray label = manager.zeros(new Shape(batchSize, 1, 5));
+                Batch batch =
+                        new Batch(
+                                manager.newSubManager(),
+                                new NDList(input),
+                                new NDList(label),
+                                batchSize,
+                                Batchifier.STACK,
+                                Batchifier.STACK,
+                                0,
+                                0);
+                PairList<String, Parameter> parameters = yolov3.getParameters();
+                EasyTrain.trainBatch(trainer, batch);
+                trainer.step();
+
+                // expected shapes
+                Shape expectedAtIndex312 =
+                        new Shape(512, 1024, 1, 1); // the layer where output0 concat to 1
+                Shape expectedAtIndex330 =
+                        new Shape(1024, 512, 3, 3); // the layer where output1 concat to 2
+
+                // test if the shape of the kernel sizes are right
+                Assert.assertEquals(
+                        parameters.get(312).getValue().getArray().getShape(), expectedAtIndex312);
+                Assert.assertEquals(
+                        parameters.get(330).getValue().getArray().getShape(), expectedAtIndex330);
+            }
+        }
+    }
+}
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/TabNetTest.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/TabNetTest.java
index 5dd6d2344..de66efd2d 100644
--- a/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/TabNetTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/TabNetTest.java
@@ -29,7 +29,7 @@ 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.training.loss.TabNetLoss;
+import ai.djl.training.loss.TabNetRegressionLoss;
 import ai.djl.translate.Batchifier;
 import ai.djl.util.PairList;
 
@@ -61,14 +61,14 @@ public class TabNetTest {
     @Test
     public void testTrainingAndLogic() {
         TrainingConfig config =
-                new DefaultTrainingConfig(new TabNetLoss())
+                new DefaultTrainingConfig(new TabNetRegressionLoss())
                         .optDevices(Engine.getInstance().getDevices(2));
 
         Block tabNet = TabNet.builder().setOutDim(10).build();
         try (Model model = Model.newInstance("tabNet")) {
             model.setBlock(tabNet);
             try (Trainer trainer = model.newTrainer(config)) {
-                int batchSize = 1;
+                int batchSize = 32;
                 Shape inputShape = new Shape(batchSize, 128);
                 trainer.initialize(inputShape);
                 NDManager manager = trainer.getManager();
@@ -89,7 +89,7 @@ public class TabNetTest {
                 trainer.step();
                 // the gamma of batchNorm Layer
                 Assert.assertEquals(
-                        parameters.get(0).getValue().getArray().getShape(), new Shape(1));
+                        parameters.get(0).getValue().getArray().getShape(), new Shape(128));
 
                 // weight of shared fullyConnected Block0
                 Assert.assertEquals(