Repository: deepjavalibrary/djl
Pull request #1688 (2022-06-01 12:21:27)

allow indexer to attach specific manager
This feature allows to pass an NDManager to NDArray indexation operation so that the result is attached to the passed NDManager instead of to the original NDManager of the NDArray.
Thus, ArrayDataset.get() does not need to attach the result of the indexation a second time, thus batch creation is speeded up.

In my performance tests, I got following improvement:

Before:
It took 6872ms totally
It took 6876ms totally
It took 6825ms totally
It took 6923ms totally

After:
It took 6734ms totally
It took 6709ms totally
It took 6728ms totally
It took 6703ms totally

So the training time is reduced by 2.2% !

Also, this pull request fixes a bug that temporary resources are returned to the wrong manager.



Patch:
diff --git a/api/src/main/java/ai/djl/Application.java b/api/src/main/java/ai/djl/Application.java
index cb65a3c66..5b5dbe8bf 100644
--- a/api/src/main/java/ai/djl/Application.java
+++ b/api/src/main/java/ai/djl/Application.java
@@ -289,4 +289,11 @@ public class Application {
          */
         Application SOFTMAX_REGRESSION = new Application("tabular/softmax_regression");
     }
+
+    /** The common set of applications for audio data. */
+    public interface Audio {
+
+        /** Any audio application, including those in {@link Audio}. */
+        Application ANY = new Application("audio");
+    }
 }
diff --git a/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java b/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java
index dea3a8c67..145621052 100644
--- a/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java
+++ b/api/src/main/java/ai/djl/modality/cv/output/Rectangle.java
@@ -89,7 +89,7 @@ public class Rectangle implements BoundingBox {
         Rectangle rec = (Rectangle) box;
         // caculate intesection lrtb
         double left = Math.max(getX(), rec.getX());
-        double top = Math.min(getY(), rec.getY());
+        double top = Math.max(getY(), rec.getY());
         double right = Math.min(getX() + getWidth(), rec.getX() + rec.getWidth());
         double bottom = Math.min(getY() + getHeight(), rec.getY() + rec.getHeight());
         double intersection = (right - left) * (bottom - top);
diff --git a/api/src/main/java/ai/djl/ndarray/NDArray.java b/api/src/main/java/ai/djl/ndarray/NDArray.java
index a48e56812..5c1ed1311 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArray.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArray.java
@@ -456,7 +456,7 @@ public interface NDArray extends NDResource, BytesSupplier {
      *     index
      */
     default void set(NDIndex index, NDArray value) {
-        getNDArrayInternal().getIndexer().set(this, index, value);
+        getNDArrayInternal().getIndexer(getManager()).set(this, index, value);
     }
 
     /**
@@ -466,7 +466,7 @@ public interface NDArray extends NDResource, BytesSupplier {
      * @param value the value to replace with
      */
     default void set(NDIndex index, Number value) {
-        getNDArrayInternal().getIndexer().set(this, index, value);
+        getNDArrayInternal().getIndexer(getManager()).set(this, index, value);
     }
 
     /**
@@ -498,7 +498,7 @@ public interface NDArray extends NDResource, BytesSupplier {
      * @throws IllegalArgumentException thrown if the index does not correspond to a single element
      */
     default void setScalar(NDIndex index, Number value) {
-        getNDArrayInternal().getIndexer().setScalar(this, index, value);
+        getNDArrayInternal().getIndexer(getManager()).setScalar(this, index, value);
     }
 
     /**
@@ -508,7 +508,18 @@ public interface NDArray extends NDResource, BytesSupplier {
      * @return the partial {@code NDArray}
      */
     default NDArray get(NDIndex index) {
-        return getNDArrayInternal().getIndexer().get(this, index);
+        return get(getManager(), index);
+    }
+
+    /**
+     * Returns a partial {@code NDArray}.
+     *
+     * @param manager the manager used to create the arrays
+     * @param index the section of this {@code NDArray} to return
+     * @return the partial {@code NDArray}
+     */
+    default NDArray get(NDManager manager, NDIndex index) {
+        return getNDArrayInternal().getIndexer(manager).get(this, index);
     }
 
     /**
@@ -549,6 +560,18 @@ public interface NDArray extends NDResource, BytesSupplier {
         return get(new NDIndex(indices));
     }
 
+    /**
+     * Returns a partial {@code NDArray}.
+     *
+     * @param manager the manager used to create the arrays
+     * @param indices the indices with each index corresponding to the dimensions and negative
+     *     indices starting from the end
+     * @return the partial {@code NDArray}
+     */
+    default NDArray get(NDManager manager, long... indices) {
+        return get(manager, new NDIndex(indices));
+    }
+
     /**
      * Returns a partial {@code NDArray} pointed by the indexed array. Given NDArray arr, NDArray
      * idx, and long axis, the output is out_{ijk} = arr_{idx_{ijk}, j, k} if axis=0 or arr_{i,
@@ -561,7 +584,8 @@ public interface NDArray extends NDResource, BytesSupplier {
     NDArray gather(NDArray index, int axis);
 
     /**
-     * Returns a partial {@code NDArray} pointed by the indexed array, according to linear indexing.
+     * Returns a partial {@code NDArray} pointed by the indexed array according to linear indexing,
+     * and the of output is of the same shape as index.
      *
      * @param index picks the elements of an NDArray to the same position as index
      * @return the partial {@code NDArray} of the same shape as index
diff --git a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
index c36762010..607fd85b7 100644
--- a/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
+++ b/api/src/main/java/ai/djl/ndarray/NDArrayAdapter.java
@@ -80,8 +80,8 @@ public abstract class NDArrayAdapter implements NDArray {
     /** {@inheritDoc} */
     @Override
     public void tempAttach(NDManager manager) {
-        detach();
         NDManager original = this.manager;
+        detach();
         this.manager = manager;
         manager.tempAttachInternal(original, getUid(), this);
     }
diff --git a/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java b/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
index 31f400812..fd29e8f8b 100644
--- a/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
+++ b/api/src/main/java/ai/djl/ndarray/index/NDArrayIndexer.java
@@ -59,7 +59,7 @@ public abstract class NDArrayIndexer {
         if (!indices.isEmpty() && indices.get(0) instanceof NDIndexBooleans) {
             if (indices.size() != 1) {
                 throw new IllegalArgumentException(
-                        "get() currently didn't support more that one boolean NDArray");
+                        "get() currently doesn't support more that one boolean NDArray");
             }
             return array.booleanMask(((NDIndexBooleans) indices.get(0)).getIndex());
         }
diff --git a/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java b/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
index 4a405308b..7aa69fd42 100644
--- a/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
+++ b/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
@@ -486,9 +486,10 @@ public interface NDArrayEx {
     /**
      * Returns an {@link NDArrayIndexer}.
      *
+     * @param manager the manager used to create the arrays
      * @return an {@link NDArrayIndexer}
      */
-    NDArrayIndexer getIndexer();
+    NDArrayIndexer getIndexer(NDManager manager);
 
     /**
      * Returns elements chosen from the {@code NDArray} or the other {@code NDArray} depending on
diff --git a/api/src/main/java/ai/djl/nn/norm/BatchNorm.java b/api/src/main/java/ai/djl/nn/norm/BatchNorm.java
index 111145267..0774e0e63 100644
--- a/api/src/main/java/ai/djl/nn/norm/BatchNorm.java
+++ b/api/src/main/java/ai/djl/nn/norm/BatchNorm.java
@@ -10,6 +10,7 @@
  * 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.nn.norm;
 
 import ai.djl.Device;
@@ -77,7 +78,7 @@ public class BatchNorm extends AbstractBlock {
     private Parameter runningMean;
     private Parameter runningVar;
 
-    BatchNorm(Builder builder) {
+    BatchNorm(BaseBuilder<?> builder) {
         super(VERSION);
         axis = builder.axis;
         epsilon = builder.epsilon;
@@ -276,20 +277,37 @@ public class BatchNorm extends AbstractBlock {
      *
      * @return a new builder
      */
-    public static Builder builder() {
+    public static BaseBuilder<?> builder() {
         return new Builder();
     }
 
     /** The Builder to construct a {@link BatchNorm}. */
-    public static final class Builder {
+    public static class Builder extends BaseBuilder<Builder> {
+        Builder() {}
 
-        private int axis = 1;
-        private float epsilon = 1E-5f;
-        private float momentum = .9f;
-        private boolean center = true;
-        private boolean scale = true;
+        /** {@inheritDoc} */
+        @Override
+        public BatchNorm build() {
+            return new BatchNorm(this);
+        }
 
-        Builder() {}
+        /** {@inheritDoc} */
+        @Override
+        public Builder self() {
+            return this;
+        }
+    }
+
+    /** The Builder to construct a {@link BatchNorm} type of {@link ai.djl.nn.Block}. */
+    public abstract static class BaseBuilder<T extends BaseBuilder<T>> {
+
+        protected int axis = 1;
+        protected float epsilon = 1E-5f;
+        protected float momentum = .9f;
+        protected boolean center = true;
+        protected boolean scale = true;
+
+        protected BaseBuilder() {}
 
         /**
          * Set the axis in which channel is specified. Defaults to 1.
@@ -297,9 +315,9 @@ public class BatchNorm extends AbstractBlock {
          * @param axis the axis in which channel is specified
          * @return this Builder
          */
-        public Builder optAxis(int axis) {
+        public T optAxis(int axis) {
             this.axis = axis;
-            return this;
+            return self();
         }
 
         /**
@@ -308,9 +326,9 @@ public class BatchNorm extends AbstractBlock {
          * @param val True or False on whether to add and train offset value
          * @return this Builder
          */
-        public Builder optCenter(boolean val) {
+        public T optCenter(boolean val) {
             center = val;
-            return this;
+            return self();
         }
 
         /**
@@ -319,9 +337,9 @@ public class BatchNorm extends AbstractBlock {
          * @param val True or False on whether to add and train scale value
          * @return this Builder
          */
-        public Builder optScale(boolean val) {
+        public T optScale(boolean val) {
             scale = val;
-            return this;
+            return self();
         }
 
         /**
@@ -330,9 +348,9 @@ public class BatchNorm extends AbstractBlock {
          * @param val the epsilon value
          * @return this Builder
          */
-        public Builder optEpsilon(float val) {
+        public T optEpsilon(float val) {
             epsilon = val;
-            return this;
+            return self();
         }
 
         /**
@@ -341,18 +359,23 @@ public class BatchNorm extends AbstractBlock {
          * @param val the momentum for moving average
          * @return this Builder
          */
-        public Builder optMomentum(float val) {
+        public T optMomentum(float val) {
             momentum = val;
-            return this;
+            return self();
         }
 
         /**
-         * Builds a {@link BatchNorm} block.
+         * Builds the new {@link BatchNorm}.
          *
-         * @return the {@link BatchNorm} block
+         * @return the new {@link BatchNorm}
          */
-        public BatchNorm build() {
-            return new BatchNorm(this);
-        }
+        public abstract BatchNorm build();
+
+        /**
+         * Returns this {code Builder} object.
+         *
+         * @return this {@code BaseBuilder}
+         */
+        public abstract T self();
     }
 }
diff --git a/api/src/main/java/ai/djl/nn/norm/GhostBatchNorm.java b/api/src/main/java/ai/djl/nn/norm/GhostBatchNorm.java
new file mode 100644
index 000000000..c2a47e65d
--- /dev/null
+++ b/api/src/main/java/ai/djl/nn/norm/GhostBatchNorm.java
@@ -0,0 +1,147 @@
+/*
+ * 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.nn.norm;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.training.ParameterStore;
+import ai.djl.translate.Batchifier;
+import ai.djl.translate.StackBatchifier;
+import ai.djl.util.PairList;
+
+/**
+ * {@link GhostBatchNorm} is similar to {@link BatchNorm} except that it splits a batch into a
+ * smaller sub-batches aka <em>ghost batches</em>, and normalize them individually to have a mean of
+ * 0 and variance of 1 and finally concatenate them again to a single batch. Each of the
+ * mini-batches contains a virtualBatchSize samples.
+ *
+ * @see <a href="https://arxiv.org/abs/1705.08741">Ghost Normalization Paper</a>
+ */
+public class GhostBatchNorm extends BatchNorm {
+
+    private int virtualBatchSize;
+    private Batchifier batchifier;
+
+    protected GhostBatchNorm(Builder builder) {
+        super(builder);
+        this.virtualBatchSize = builder.virtualBatchSize;
+        this.batchifier = new StackBatchifier();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected NDList forwardInternal(
+            ParameterStore parameterStore,
+            NDList inputs,
+            boolean training,
+            PairList<String, Object> params) {
+
+        NDList[] subBatches = split(inputs);
+
+        for (int i = 0; i < subBatches.length; i++) {
+            subBatches[i] = super.forwardInternal(parameterStore, subBatches[i], training, params);
+        }
+
+        return batchify(subBatches);
+    }
+
+    /**
+     * Splits an {@link NDList} into the given <b>size</b> of sub-batch.
+     *
+     * <p>This function unbatchifies the input {@link NDList} into mini-batches, each with the size
+     * of virtualBatchSize. If the batch size is divisible by the virtual batch size, all returned
+     * sub-batches will be the same size. If the batch size is not divisible by virtual batch size,
+     * all returned sub-batches will be the same size, except the last one.
+     *
+     * @param list the {@link NDList} that needs to be split
+     * @return an array of {@link NDList} that contains all the mini-batches
+     */
+    protected NDList[] split(NDList list) {
+        double batchSize = list.head().size(0);
+        int countBatches = (int) Math.ceil(batchSize / virtualBatchSize);
+
+        return batchifier.split(list, countBatches, true);
+    }
+
+    /**
+     * Converts an array of {@link NDList} into an NDList using {@link StackBatchifier} and squeezes
+     * the first dimension created by it. This makes the final {@link NDArray} same size as the
+     * splitted one.
+     *
+     * @param subBatches the input array of {@link NDList}
+     * @return the batchified {@link NDList}
+     */
+    protected NDList batchify(NDList[] subBatches) {
+        NDList batch = batchifier.batchify(subBatches);
+
+        return squeezeExtraDimensions(batch);
+    }
+
+    /**
+     * Squeezes first axes of {@link NDList}.
+     *
+     * @param batch input array of {@link NDList}
+     * @return the squeezed {@link NDList}
+     */
+    protected NDList squeezeExtraDimensions(NDList batch) {
+        NDArray array = batch.singletonOrThrow().squeeze(0);
+        batch.set(0, array);
+
+        return batch;
+    }
+
+    /**
+     * Creates a builder to build a {@code GhostBatchNorm}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** The Builder to construct a {@link GhostBatchNorm}. */
+    public static class Builder extends BatchNorm.BaseBuilder<Builder> {
+
+        private int virtualBatchSize = 128;
+
+        Builder() {}
+
+        /**
+         * Sets the size of virtual batches in which to use when sub-batching. Defaults to 128.
+         *
+         * @param virtualBatchSize the virtual batch size
+         * @return this Builder
+         */
+        public Builder optVirtualBatchSize(int virtualBatchSize) {
+            this.virtualBatchSize = virtualBatchSize;
+            return this;
+        }
+
+        /**
+         * Builds the new {@link GhostBatchNorm}.
+         *
+         * @return the new {@link GhostBatchNorm}
+         */
+        @Override
+        public GhostBatchNorm build() {
+            return new GhostBatchNorm(this);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Builder self() {
+            return this;
+        }
+    }
+}
diff --git a/api/src/main/java/ai/djl/training/dataset/ArrayDataset.java b/api/src/main/java/ai/djl/training/dataset/ArrayDataset.java
index 997d70527..57f0653c3 100644
--- a/api/src/main/java/ai/djl/training/dataset/ArrayDataset.java
+++ b/api/src/main/java/ai/djl/training/dataset/ArrayDataset.java
@@ -21,19 +21,27 @@ import java.util.stream.Stream;
 
 /**
  * {@code ArrayDataset} is an implementation of {@link RandomAccessDataset} that consist entirely of
- * large {@link NDArray}s. There can be multiple data and label {@link NDArray}s within the dataset.
- * Each sample will be retrieved by indexing each {@link NDArray} along the first dimension.
+ * large {@link NDArray}s. It is recommended only for datasets small enough to fit in memory that
+ * come in array formats. Otherwise, consider directly using the {@link RandomAccessDataset}
+ * instead.
+ *
+ * <p>There can be multiple data and label {@link NDArray}s within the dataset. Each sample will be
+ * retrieved by indexing each {@link NDArray} along the first dimension.
  *
  * <p>The following is an example of how to use ArrayDataset:
  *
  * <pre>
  *     ArrayDataset dataset = new ArrayDataset.Builder()
- *                              .setData(data)
- *                              .optLabels(label)
+ *                              .setData(data1, data2)
+ *                              .optLabels(labels1, labels2, labels3)
  *                              .setSampling(20, false)
  *                              .build();
  * </pre>
  *
+ * <p>Suppose you get a {@link Batch} from {@code trainer.iterateDataset(dataset)} or {@code
+ * dataset.getData(manager)}. In the data of this batch, it will be an NDList with one NDArray for
+ * each data input. In this case, it would be 2 arrays. Similarly, the labels would have 3 arrays.
+ *
  * @see Dataset
  */
 public class ArrayDataset extends RandomAccessDataset {
@@ -76,15 +84,13 @@ public class ArrayDataset extends RandomAccessDataset {
         NDList datum = new NDList();
         NDList label = new NDList();
         for (NDArray array : data) {
-            datum.add(array.get(index));
+            datum.add(array.get(manager, index));
         }
         if (labels != null) {
             for (NDArray array : labels) {
-                label.add(array.get(index));
+                label.add(array.get(manager, index));
             }
         }
-        datum.attach(manager);
-        label.attach(manager);
         return new Record(datum, label);
     }
 
diff --git a/api/src/main/java/ai/djl/training/dataset/Dataset.java b/api/src/main/java/ai/djl/training/dataset/Dataset.java
index 90a21da79..9dacf1694 100644
--- a/api/src/main/java/ai/djl/training/dataset/Dataset.java
+++ b/api/src/main/java/ai/djl/training/dataset/Dataset.java
@@ -21,7 +21,9 @@ import java.util.concurrent.ExecutorService;
 /**
  * An interface to represent a set of sample data/label pairs to train a model.
  *
- * @see <a href="http://docs.djl.ai/docs/dataset.html">The guide on datasets</a>
+ * @see <a href="http://docs.djl.ai/docs/dataset.html">The guide to datasets</a>
+ * @see <a href="http://docs.djl.ai/docs/development/how_to_use_dataset.html">The guide to
+ *     implementing a custom dataset</a>
  */
 public interface Dataset {
 
diff --git a/api/src/main/java/ai/djl/training/dataset/RandomAccessDataset.java b/api/src/main/java/ai/djl/training/dataset/RandomAccessDataset.java
index eab783fec..0d5e20623 100644
--- a/api/src/main/java/ai/djl/training/dataset/RandomAccessDataset.java
+++ b/api/src/main/java/ai/djl/training/dataset/RandomAccessDataset.java
@@ -34,6 +34,11 @@ import java.util.stream.IntStream;
 /**
  * RandomAccessDataset represent the dataset that support random access reads. i.e. it could access
  * a specific data item given the index.
+ *
+ * <p>Almost all datasets in DJL extend, either directly or indirectly, {@link RandomAccessDataset}.
+ *
+ * @see <a href="http://docs.djl.ai/docs/development/how_to_use_dataset.html">The guide to
+ *     implementing a custom dataset</a>
  */
 public abstract class RandomAccessDataset implements Dataset {
 
diff --git a/api/src/main/java/ai/djl/util/Utils.java b/api/src/main/java/ai/djl/util/Utils.java
index 6feb8fac5..2fc5d7eef 100644
--- a/api/src/main/java/ai/djl/util/Utils.java
+++ b/api/src/main/java/ai/djl/util/Utils.java
@@ -19,6 +19,7 @@ import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
+import java.nio.file.FileVisitOption;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
@@ -243,7 +244,7 @@ public final class Utils {
     public static int getCurrentEpoch(Path modelDir, String modelName) throws IOException {
         final Pattern pattern = Pattern.compile(Pattern.quote(modelName) + "-(\\d{4}).params");
         List<Integer> checkpoints =
-                Files.walk(modelDir, 1)
+                Files.walk(modelDir, 1, FileVisitOption.FOLLOW_LINKS)
                         .map(
                                 p -> {
                                     Matcher m = pattern.matcher(p.toFile().getName());
diff --git a/basicdataset/README.md b/basicdataset/README.md
index c342cd15d..4cd0afbc8 100644
--- a/basicdataset/README.md
+++ b/basicdataset/README.md
@@ -4,17 +4,7 @@
 
 This module contains a number of basic and standard datasets in the Deep Java Library's (DJL). These datasets are used to train deep learning models.
 
-## List of datasets
-
-This module contains the following datasets:
-
-- [MNIST](http://yann.lecun.com/exdb/mnist/) - A handwritten digits dataset
-- [CIFAR10](https://www.cs.toronto.edu/~kriz/cifar.html) - A dataset consisting of 60,000 32x32 color images in 10 classes
-- [Coco](http://cocodataset.org) - A large-scale object detection, segmentation, and captioning dataset that contains 1.5 million object instances
-    - You have to manually add `com.twelvemonkeys.imageio:imageio-jpeg:3.5` dependency to your project
-- [ImageNet](http://www.image-net.org/) - An image database organized according to the WordNet hierarchy
-  >**Note**: You have to manually download the ImageNet dataset due to licensing requirements.
-- [Pikachu](http://d2l.ai/chapter_computer-vision/object-detection-dataset.html) - 1000 Pikachu images of different angles and sizes created using an open source 3D Pikachu model
+You can find the datasets provided by this module on our [docs](http://docs.djl.ai/docs/dataset.html).
 
 ## Documentation
 
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 3cdb4c022..b36f2843d 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/CocoDetection.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/CocoDetection.java
@@ -34,6 +34,13 @@ import java.util.Optional;
 /**
  * Coco image detection dataset from http://cocodataset.org/#home.
  *
+ * <p>Coco is a large-scale object detection, segmentation, and captioning dataset although only
+ * object detection is implemented at thsi time. It contains 1.5 million object instances and is one
+ * of the standard benchmark object detection datasets.
+ *
+ * <p>To use this dataset, you have to manually add {@code
+ * com.twelvemonkeys.imageio:imageio-jpeg:3.5} as a dependency in your project.
+ *
  * <p>Each image might have different {@link ai.djl.ndarray.types.Shape}s.
  */
 public class CocoDetection extends ObjectDetectionDataset {
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 0a7130dfa..f1ad1f387 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/PikachuDetection.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/PikachuDetection.java
@@ -39,7 +39,14 @@ import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
-/** Pikachu image detection dataset that contains multiple Pikachus in each image. */
+/**
+ * Pikachu image detection dataset that contains multiple Pikachus in each image.
+ *
+ * <p>It was based on a section from the [Dive into Deep Learning
+ * book](http://d2l.ai/chapter_computer-vision/object-detection-dataset.html). It contains 1000
+ * Pikachu images of different angles and sizes created using an open source 3D Pikachu model. Each
+ * image contains only a single pikachu.
+ */
 public class PikachuDetection extends ObjectDetectionDataset {
 
     private static final String VERSION = "1.0";
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/AbstractImageFolder.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/AbstractImageFolder.java
index fe8ccad10..b370ad5d6 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/AbstractImageFolder.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/AbstractImageFolder.java
@@ -34,7 +34,11 @@ import java.util.stream.Stream;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
-/** A dataset for loading image files stored in a folder structure. */
+/**
+ * A dataset for loading image files stored in a folder structure.
+ *
+ * <p>Usually, you want to use {@link ImageFolder} instead.
+ */
 public abstract class AbstractImageFolder extends ImageClassificationDataset {
 
     private static final Logger logger = LoggerFactory.getLogger(AbstractImageFolder.class);
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Cifar10.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Cifar10.java
index 21d7a5bc0..5e9402728 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Cifar10.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Cifar10.java
@@ -34,6 +34,9 @@ import java.util.Map;
 /**
  * CIFAR10 image classification dataset from https://www.cs.toronto.edu/~kriz/cifar.html.
  *
+ * <p>It consists of 60,000 32x32 color images with 10 classes. It can train in a few hours with a
+ * GPU.
+ *
  * <p>Each sample is an image (in 3-D {@link NDArray}) with shape (32, 32, 3).
  */
 public final class Cifar10 extends ArrayDataset {
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/FashionMnist.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/FashionMnist.java
index abbad625b..90f1c0e44 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/FashionMnist.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/FashionMnist.java
@@ -33,9 +33,11 @@ import java.util.Map;
 
 /**
  * FashMnist is a dataset from Zalando article images
- * https://github.com/zalandoresearch/fashion-mnist.
+ * (https://github.com/zalandoresearch/fashion-mnist).
  *
- * <p>Each sample is an image (in 3-D NDArray) with shape (28, 28, 1).
+ * <p>Each sample is a grayscale image (in 3-D NDArray) with shape (28, 28, 1).
+ *
+ * <p>It was created to be a drop in replacement for {@link Mnist}, but have a less simplistic task.
  */
 public final class FashionMnist extends ArrayDataset {
 
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/ImageFolder.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/ImageFolder.java
index fa2251477..b9da9ee02 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/ImageFolder.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/ImageFolder.java
@@ -24,6 +24,8 @@ import java.util.Arrays;
 /**
  * A dataset for loading image files stored in a folder structure.
  *
+ * <p>Below is an example directory layout for the image folder:
+ *
  * <pre>
  *  The image folder should be structured as follows:
  *       root/shoes/Aerobic Shoes1.png
@@ -32,11 +34,35 @@ import java.util.Arrays;
  *       root/boots/Black Boots.png
  *       root/boots/White Boots.png
  *       ...
- *       root/pumps/Red Pumps
- *       root/pumps/Pink Pumps
+ *       root/pumps/Red Pumps.png
+ *       root/pumps/Pink Pumps.png
  *       ...
+ *
  *  here shoes, boots, pumps are your labels
  *  </pre>
+ *
+ * <p>Here, the dataset will take the folder names (shoes, boots, bumps) in sorted order as your
+ * labels. Nested folder structures are not currently supported.
+ *
+ * <p>Then, you can create your instance of the dataset as follows:
+ *
+ * <pre>
+ * // set the image folder path
+ * Repository repository = Repository.newInstance("folder", Paths.get("/path/to/imagefolder/root");
+ * ImageFolder dataset =
+ *     new ImageFolder.Builder()
+ *         .setRepository(repository)
+ *         .addTransform(new Resize(100, 100)) // Use image transforms as necessary for your data
+ *         .addTransform(new ToTensor()) // Usually required as the last transform to convert images to tensors
+ *         .setSampling(batchSize, true)
+ *         .build();
+ *
+ * // call prepare before using
+ * dataset.prepare();
+ *
+ * // to get the synset or label names
+ * List&gt;String&lt; synset = dataset.getSynset();
+ * </pre>
  */
 public final class ImageFolder extends AbstractImageFolder {
 
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Mnist.java b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Mnist.java
index 257fd6bc4..1cb87668d 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Mnist.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Mnist.java
@@ -34,7 +34,11 @@ import java.util.Map;
 /**
  * MNIST handwritten digits dataset from http://yann.lecun.com/exdb/mnist.
  *
- * <p>Each sample is an image (in 3-D NDArray) with shape (28, 28, 1).
+ * <p>Each sample is a grayscale image (in 3-D NDArray) with shape (28, 28, 1).
+ *
+ * <p>It is a common starting dataset because it is small and can train within minutes. However, it
+ * is an overly easy task that even poor models can still perform very well on. Instead, consider
+ * {@link FashionMnist} which offers a comparable speed but a more reasonable difficulty task.
  */
 public final class Mnist extends ArrayDataset {
 
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/nlp/StanfordQuestionAnsweringDataset.java b/basicdataset/src/main/java/ai/djl/basicdataset/nlp/StanfordQuestionAnsweringDataset.java
index 7716d5ab1..aa3d0280d 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/nlp/StanfordQuestionAnsweringDataset.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/nlp/StanfordQuestionAnsweringDataset.java
@@ -39,6 +39,8 @@ import java.util.Map;
  * questions posed by crowdworkers on a set of Wikipedia articles, where the answer to every
  * question is a segment of text, or span, from the corresponding reading passage, or the question
  * might be unanswerable.
+ *
+ * @see <a href="https://rajpurkar.github.io/SQuAD-explorer/">Dataset website</a>
  */
 @SuppressWarnings("unchecked")
 public class StanfordQuestionAnsweringDataset extends TextDataset implements RawDataset<Object> {
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/nlp/UniversalDependenciesEnglishEWT.java b/basicdataset/src/main/java/ai/djl/basicdataset/nlp/UniversalDependenciesEnglishEWT.java
new file mode 100644
index 000000000..62a5627f5
--- /dev/null
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/nlp/UniversalDependenciesEnglishEWT.java
@@ -0,0 +1,219 @@
+/*
+ * 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.basicdataset.nlp;
+
+import ai.djl.Application.NLP;
+import ai.djl.basicdataset.BasicDatasets;
+import ai.djl.modality.nlp.embedding.EmbeddingException;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
+import ai.djl.repository.Artifact;
+import ai.djl.repository.MRL;
+import ai.djl.training.dataset.Record;
+import ai.djl.util.Progress;
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * A Gold Standard Universal Dependencies Corpus for English, built over the source material of the
+ * English Web Treebank LDC2012T13.
+ *
+ * @see <a href="https://catalog.ldc.upenn.edu/LDC2012T13">English Web Treebank LDC2012T13</a>
+ */
+public class UniversalDependenciesEnglishEWT extends TextDataset {
+
+    private static final String VERSION = "2.0";
+    private static final String ARTIFACT_ID = "universal-dependencies-en-ewt";
+
+    private List<List<Integer>> universalPosTags;
+
+    /**
+     * Creates a new instance of {@code UniversalDependenciesEnglish}.
+     *
+     * @param builder the builder object to build from
+     */
+    protected UniversalDependenciesEnglishEWT(Builder builder) {
+        super(builder);
+        this.usage = builder.usage;
+        mrl = builder.getMrl();
+    }
+
+    /**
+     * Creates a new builder to build a {@link UniversalDependenciesEnglishEWT}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /**
+     * Prepares the dataset for use with tracked progress. In this method the TXT file will be
+     * parsed. The texts will be added to {@code sourceTextData} and the Universal POS tags will be
+     * added to {@code universalPosTags}. Only {@code sourceTextData} will then be preprocessed.
+     *
+     * @param progress the progress tracker
+     * @throws IOException for various exceptions depending on the dataset
+     * @throws EmbeddingException if there are exceptions during the embedding process
+     */
+    @Override
+    public void prepare(Progress progress) throws IOException, EmbeddingException {
+        if (prepared) {
+            return;
+        }
+
+        Artifact artifact = mrl.getDefaultArtifact();
+        mrl.prepare(artifact, progress);
+        Path root = mrl.getRepository().getResourceDirectory(artifact);
+        Path usagePath = null;
+        switch (usage) {
+            case TRAIN:
+                usagePath = Paths.get("en-ud-v2/en-ud-v2/en-ud-tag.v2.train.txt");
+                break;
+            case TEST:
+                usagePath = Paths.get("en-ud-v2/en-ud-v2/en-ud-tag.v2.test.txt");
+                break;
+            case VALIDATION:
+                usagePath = Paths.get("en-ud-v2/en-ud-v2/en-ud-tag.v2.dev.txt");
+                break;
+            default:
+                break;
+        }
+        usagePath = root.resolve(usagePath);
+
+        StringBuilder sourceTextDatum = new StringBuilder();
+        List<String> sourceTextData = new ArrayList<>();
+        universalPosTags = new ArrayList<>();
+        List<Integer> universalPosTag = new ArrayList<>();
+        try (BufferedReader reader = Files.newBufferedReader(usagePath)) {
+            String row;
+            while ((row = reader.readLine()) != null) {
+                if (("").equals(row)) {
+                    sourceTextData.add(sourceTextDatum.toString());
+                    universalPosTags.add(universalPosTag);
+
+                    sourceTextDatum.delete(0, sourceTextDatum.length());
+                    universalPosTag = new ArrayList<>();
+
+                    continue;
+                }
+                String[] splits = row.split("\t");
+                if (sourceTextDatum.length() != 0) {
+                    sourceTextDatum.append(' ');
+                }
+                sourceTextDatum.append(splits[0]);
+                universalPosTag.add(UniversalPosTag.valueOf(splits[1]).ordinal());
+            }
+        }
+
+        preprocess(sourceTextData, true);
+        prepared = true;
+    }
+
+    /**
+     * Gets the {@link Record} for the given index from the dataset.
+     *
+     * @param manager the manager used to create the arrays
+     * @param index the index of the requested data item
+     * @return a {@link Record} that contains the data and label of the requested data item. The
+     *     data {@link NDList} contains one {@link NDArray} representing the text embedding, The
+     *     label {@link NDList} contains one {@link NDArray} including the indices of the Universal
+     *     POS tags of each token. For the index of each Universal POS tag, see the enum class
+     *     {@link UniversalPosTag}.
+     */
+    @Override
+    public Record get(NDManager manager, long index) {
+        NDList data = new NDList(sourceTextData.getEmbedding(manager, index));
+        NDList labels =
+                new NDList(
+                        manager.create(
+                                        universalPosTags
+                                                .get(Math.toIntExact(index))
+                                                .stream()
+                                                .mapToInt(Integer::intValue)
+                                                .toArray())
+                                .toType(DataType.INT32, false));
+        return new Record(data, labels);
+    }
+
+    /**
+     * Returns the number of records available to be read in this {@code Dataset}.
+     *
+     * @return the number of records available to be read in this {@code Dataset}
+     */
+    @Override
+    protected long availableSize() {
+        return sourceTextData.getSize();
+    }
+
+    /** A builder for a {@link UniversalDependenciesEnglishEWT}. */
+    public static class Builder extends TextDataset.Builder<Builder> {
+
+        /** Constructs a new builder. */
+        public Builder() {
+            groupId = BasicDatasets.GROUP_ID + ".universal-dependencies";
+            artifactId = ARTIFACT_ID;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public Builder self() {
+            return this;
+        }
+
+        /**
+         * Builds the {@link UniversalDependenciesEnglishEWT}.
+         *
+         * @return the {@link UniversalDependenciesEnglishEWT}
+         */
+        public UniversalDependenciesEnglishEWT build() {
+            return new UniversalDependenciesEnglishEWT(this);
+        }
+
+        MRL getMrl() {
+            return repository.dataset(NLP.ANY, groupId, artifactId, VERSION);
+        }
+    }
+
+    /**
+     * An enum class for Universal POS tags which mark the core part-of-speech categories.
+     *
+     * @see <a href="https://universaldependencies.org/u/pos/">Universal POS tags</a>
+     */
+    enum UniversalPosTag {
+        ADJ,
+        ADV,
+        INTJ,
+        NOUN,
+        PROPN,
+        VERB,
+        ADP,
+        AUX,
+        CCONJ,
+        DET,
+        NUM,
+        PART,
+        PRON,
+        SCONJ,
+        PUNCT,
+        SYM,
+        X;
+    }
+}
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/AirfoilRandomAccess.java b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/AirfoilRandomAccess.java
index 49cfe4016..bd4d6e209 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/AirfoilRandomAccess.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/AirfoilRandomAccess.java
@@ -14,6 +14,7 @@ package ai.djl.basicdataset.tabular;
 
 import ai.djl.Application.Tabular;
 import ai.djl.basicdataset.BasicDatasets;
+import ai.djl.basicdataset.utils.Feature;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
@@ -33,7 +34,8 @@ import org.apache.commons.csv.CSVFormat;
 import org.apache.commons.csv.CSVRecord;
 
 /**
- * Airfoil Self-Noise Data Set from https://archive.ics.uci.edu/ml/datasets/Airfoil+Self-Noise.
+ * Airfoil Self-Noise Data Set from <a
+ * href="https://archive.ics.uci.edu/ml/datasets/Airfoil+Self-Noise">https://archive.ics.uci.edu/ml/datasets/Airfoil+Self-Noise</a>.
  *
  * <p>1503 instances 6 attributes
  */
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 177672d7c..a09b0e82c 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/CsvDataset.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/CsvDataset.java
@@ -13,6 +13,7 @@
 package ai.djl.basicdataset.tabular;
 
 import ai.djl.basicdataset.utils.DynamicBuffer;
+import ai.djl.basicdataset.utils.Feature;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.Shape;
@@ -31,7 +32,6 @@ import java.nio.charset.StandardCharsets;
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Collections;
-import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.zip.GZIPInputStream;
@@ -42,8 +42,6 @@ import org.apache.commons.csv.CSVRecord;
 /** {@code CsvDataset} represents the dataset that stored in a .csv file. */
 public class CsvDataset extends RandomAccessDataset {
 
-    private static final Featurizer NUMERIC_FEATURIZER = new NumericFeaturizer();
-
     protected URL csvUrl;
     protected CSVFormat csvFormat;
     protected List<Feature> features;
@@ -116,7 +114,7 @@ public class CsvDataset extends RandomAccessDataset {
         for (Feature feature : selected) {
             String name = feature.getName();
             String value = record.get(name);
-            feature.featurizer.featurize(bb, value);
+            feature.getFeaturizer().featurize(bb, value);
         }
         FloatBuffer buf = bb.getBuffer();
         return new NDList(manager.create(buf, new Shape(bb.getLength())));
@@ -291,129 +289,4 @@ public class CsvDataset extends RandomAccessDataset {
             return new CsvDataset(this);
         }
     }
-
-    /** An interface that convert String to numeric data. */
-    public interface Featurizer {
-
-        /**
-         * Puts encoded data into the float buffer.
-         *
-         * @param buf the float buffer to be filled
-         * @param input the string input
-         */
-        void featurize(DynamicBuffer buf, String input);
-    }
-
-    /** A class contains feature name and its {@code Featurizer}. */
-    public static final class Feature {
-
-        String name;
-        Featurizer featurizer;
-
-        /**
-         * Constructs a {@code Feature} instance.
-         *
-         * @param name the feature name
-         * @param featurizer the {@code Featurizer}
-         */
-        public Feature(String name, Featurizer featurizer) {
-            this.name = name;
-            this.featurizer = featurizer;
-        }
-
-        /**
-         * Constructs a {@code Feature} instance.
-         *
-         * @param name the feature name
-         * @param numeric true if input is numeric data
-         */
-        public Feature(String name, boolean numeric) {
-            this.name = name;
-            if (numeric) {
-                featurizer = NUMERIC_FEATURIZER;
-            } else {
-                featurizer = new StringFeaturizer();
-            }
-        }
-
-        /**
-         * Constructs a {@code Feature} instance.
-         *
-         * @param name the feature name
-         * @param map a map contains categorical value maps to index
-         * @param onehotEncode true if use onehot encode
-         */
-        public Feature(String name, Map<String, Integer> map, boolean onehotEncode) {
-            this.name = name;
-            this.featurizer = new StringFeaturizer(map, onehotEncode);
-        }
-
-        /**
-         * Returns the feature name.
-         *
-         * @return the feature name
-         */
-        public String getName() {
-            return name;
-        }
-
-        /**
-         * Returns the {@code Featurizer}.
-         *
-         * @return the {@code Featurizer}
-         */
-        public Featurizer getFeaturizer() {
-            return featurizer;
-        }
-    }
-
-    private static final class NumericFeaturizer implements Featurizer {
-
-        /** {@inheritDoc} */
-        @Override
-        public void featurize(DynamicBuffer buf, String input) {
-            buf.put(Float.parseFloat(input));
-        }
-    }
-
-    private static final class StringFeaturizer implements Featurizer {
-
-        private Map<String, Integer> map;
-        private boolean onehotEncode;
-        private boolean autoMap;
-
-        StringFeaturizer() {
-            this.map = new HashMap<>();
-            this.autoMap = true;
-        }
-
-        StringFeaturizer(Map<String, Integer> map, boolean onehotEncode) {
-            this.map = map;
-            this.onehotEncode = onehotEncode;
-        }
-
-        /** {@inheritDoc} */
-        @Override
-        public void featurize(DynamicBuffer buf, String input) {
-            if (onehotEncode) {
-                for (int i = 0; i < map.size(); ++i) {
-                    buf.put(i == map.get(input) ? 1 : 0);
-                }
-                return;
-            }
-
-            Integer index = map.get(input);
-            if (index != null) {
-                buf.put(index);
-                return;
-            }
-
-            if (!autoMap) {
-                throw new IllegalArgumentException("Value: " + input + " not found in the map.");
-            }
-            int value = map.size();
-            map.put(input, value);
-            buf.put(value);
-        }
-    }
 }
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/DailyDelhiClimate.java b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/DailyDelhiClimate.java
index bb27640ca..54cf4e244 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/DailyDelhiClimate.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/DailyDelhiClimate.java
@@ -15,6 +15,8 @@ package ai.djl.basicdataset.tabular;
 import ai.djl.Application.Tabular;
 import ai.djl.basicdataset.BasicDatasets;
 import ai.djl.basicdataset.utils.DynamicBuffer;
+import ai.djl.basicdataset.utils.Feature;
+import ai.djl.basicdataset.utils.Featurizer;
 import ai.djl.ndarray.NDList;
 import ai.djl.ndarray.NDManager;
 import ai.djl.repository.Artifact;
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/utils/Feature.java b/basicdataset/src/main/java/ai/djl/basicdataset/utils/Feature.java
new file mode 100644
index 000000000..46b91fd16
--- /dev/null
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/utils/Feature.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.basicdataset.utils;
+
+import java.util.Map;
+
+/** A class contains feature name and its {@code Featurizer}. */
+public final class Feature {
+
+    String name;
+    Featurizer featurizer;
+
+    /**
+     * Constructs a {@code Feature} instance.
+     *
+     * @param name the feature name
+     * @param featurizer the {@code Featurizer}
+     */
+    public Feature(String name, Featurizer featurizer) {
+        this.name = name;
+        this.featurizer = featurizer;
+    }
+
+    /**
+     * Constructs a {@code Feature} instance.
+     *
+     * @param name the feature name
+     * @param numeric true if input is numeric data
+     */
+    public Feature(String name, boolean numeric) {
+        this.name = name;
+        if (numeric) {
+            featurizer = Featurizers.getNumericFeaturizer();
+        } else {
+            featurizer = Featurizers.getStringFeaturizer();
+        }
+    }
+
+    /**
+     * Constructs a {@code Feature} instance.
+     *
+     * @param name the feature name
+     * @param map a map contains categorical value maps to index
+     * @param onehotEncode true if use onehot encode
+     */
+    public Feature(String name, Map<String, Integer> map, boolean onehotEncode) {
+        this.name = name;
+        this.featurizer = Featurizers.getStringFeaturizer(map, onehotEncode);
+    }
+
+    /**
+     * Returns the feature name.
+     *
+     * @return the feature name
+     */
+    public String getName() {
+        return name;
+    }
+
+    /**
+     * Returns the {@code Featurizer}.
+     *
+     * @return the {@code Featurizer}
+     */
+    public Featurizer getFeaturizer() {
+        return featurizer;
+    }
+}
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/utils/Featurizer.java b/basicdataset/src/main/java/ai/djl/basicdataset/utils/Featurizer.java
new file mode 100644
index 000000000..a6c8d0517
--- /dev/null
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/utils/Featurizer.java
@@ -0,0 +1,25 @@
+/*
+ * 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.basicdataset.utils;
+
+/** An interface that convert String to numeric data. */
+public interface Featurizer {
+
+    /**
+     * Puts encoded data into the float buffer.
+     *
+     * @param buf the float buffer to be filled
+     * @param input the string input
+     */
+    void featurize(DynamicBuffer buf, String input);
+}
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/utils/Featurizers.java b/basicdataset/src/main/java/ai/djl/basicdataset/utils/Featurizers.java
new file mode 100644
index 000000000..1f0d76a96
--- /dev/null
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/utils/Featurizers.java
@@ -0,0 +1,104 @@
+/*
+ * 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.basicdataset.utils;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/** A utility class provides helper functions to create {@link Featurizer}. */
+public final class Featurizers {
+
+    private static final Featurizer NUMERIC_FEATURIZER = new NumericFeaturizer();
+    private static final Featurizer STRING_FEATURIZER = new StringFeaturizer();
+
+    private Featurizers() {}
+
+    /**
+     * Returns the default numeric {@link Featurizer}.
+     *
+     * @return the default numeric {@link Featurizer}
+     */
+    public static Featurizer getNumericFeaturizer() {
+        return NUMERIC_FEATURIZER;
+    }
+
+    /**
+     * Returns the default String {@link Featurizer}.
+     *
+     * @return the default String {@link Featurizer}
+     */
+    public static Featurizer getStringFeaturizer() {
+        return STRING_FEATURIZER;
+    }
+
+    /**
+     * Returns a new instance of String {@link Featurizer}.
+     *
+     * @param map a map contains categorical value maps to index
+     * @param onehotEncode true if use onehot encode
+     * @return a new instance of String {@link Featurizer}
+     */
+    public static Featurizer getStringFeaturizer(Map<String, Integer> map, boolean onehotEncode) {
+        return new StringFeaturizer(map, onehotEncode);
+    }
+
+    private static final class NumericFeaturizer implements Featurizer {
+
+        /** {@inheritDoc} */
+        @Override
+        public void featurize(DynamicBuffer buf, String input) {
+            buf.put(Float.parseFloat(input));
+        }
+    }
+
+    private static final class StringFeaturizer implements Featurizer {
+
+        private Map<String, Integer> map;
+        private boolean onehotEncode;
+        private boolean autoMap;
+
+        StringFeaturizer() {
+            this.map = new HashMap<>();
+            this.autoMap = true;
+        }
+
+        StringFeaturizer(Map<String, Integer> map, boolean onehotEncode) {
+            this.map = map;
+            this.onehotEncode = onehotEncode;
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        public void featurize(DynamicBuffer buf, String input) {
+            if (onehotEncode) {
+                for (int i = 0; i < map.size(); ++i) {
+                    buf.put(i == map.get(input) ? 1 : 0);
+                }
+                return;
+            }
+
+            Integer index = map.get(input);
+            if (index != null) {
+                buf.put(index);
+                return;
+            }
+
+            if (!autoMap) {
+                throw new IllegalArgumentException("Value: " + input + " not found in the map.");
+            }
+            int value = map.size();
+            map.put(input, value);
+            buf.put(value);
+        }
+    }
+}
diff --git a/docs/dataset.md b/docs/dataset.md
index 5bbddeb20..cb6d10778 100644
--- a/docs/dataset.md
+++ b/docs/dataset.md
@@ -1,6 +1,6 @@
 # Dataset
 
-A dataset (or data set) is a collection of data that are used for machine-learning training job.
+A dataset (or data set) is a collection of data that is used for training a machine learning model.
 
 Machine learning typically works with three datasets:
 
@@ -11,7 +11,7 @@ Machine learning typically works with three datasets:
 - Validation dataset
 
     The validation set is used to evaluate a given model during the training process. It helps machine learning
-    engineers to fine-tune the [HyperParameter](https://github.com/deepjavalibrary/djl/blob/master/api/src/main/java/ai/djl/training/hyperparameter/param/Hyperparameter.java)
+    engineers to fine-tune the [HyperParameters](https://github.com/deepjavalibrary/djl/blob/master/api/src/main/java/ai/djl/training/hyperparameter/param/Hyperparameter.java)
     at model development stage.
     The model doesn't learn from validation dataset; and validation dataset is optional.
     
@@ -19,9 +19,59 @@ Machine learning typically works with three datasets:
 
     The Test dataset provides the gold standard used to evaluate the model.
     It is only used once a model is completely trained.
+    The test dataset should more accurately evaluate how the model will be performed on new data.
  
 See [Jason Brownlee’s article](https://machinelearningmastery.com/difference-test-validation-datasets/) for more detail.
  
 ## [Basic Dataset](../basicdataset/README.md)
  
 DJL provides a number of built-in basic and standard datasets. These datasets are used to train deep learning models.
+This module contains the following datasets:
+
+### CV
+
+#### Image Classification
+
+- [MNIST](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/classification/Mnist.html) - A small and fast handwritten digits dataset
+- [Fashion MNIST](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/classification/FashionMnist.html) - A small and fast clothing type detection dataset
+- [CIFAR10](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/classification/Cifar10.html) - A dataset consisting of 60,000 32x32 color images in 10 classes
+- [ImageNet](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/classification/ImageNet.html) - An image database organized according to the WordNet hierarchy
+  >**Note**: You have to manually download the ImageNet dataset due to licensing requirements.
+
+#### Object Detection
+
+- [Pikachu](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/PikachuDetection.html) - 1000 Pikachu images of different angles and sizes created using an open source 3D Pikachu model
+- [Banana Detection](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/BananaDetection.html) - A testing single object detection dataset
+
+#### Other CV
+
+- [Captcha](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/classification/CaptchaDataset.html) - A dataset for a grayscale 6-digit CAPTCHA task
+- [Coco](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/CocoDetection.html) - A large-scale object detection, segmentation, and captioning dataset that contains 1.5 million object instances
+  - You have to manually add `com.twelvemonkeys.imageio:imageio-jpeg:3.5` dependency to your project
+
+### NLP
+
+#### Text Classification and Sentiment Analysis
+
+- [AmazonReview](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/nlp/AmazonReview.html) - A sentiment analysis dataset of Amazon Reviews with their ratings
+- [Stanford Movie Review](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/nlp/StanfordMovieReview.html) - A sentiment analysis dataset of movie reviews and sentiments sourced from IMDB
+- [GoEmotions](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/nlp/GoEmotions.html) - A dataset classifying 50k curated reddit comments into either 27 emotion categories or neutral
+
+#### Unlabeled Text
+
+- [Penn Treebank Text](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/nlp/PennTreebankText.html) - The text (not POS tags) from the Penn Treebank, a collection of Wall Street Journal stories
+- [WikiText2](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/nlp/WikiText2.html) - A collection of over 100 million tokens extracted from good and featured articles on wikipedia
+ 
+#### Other NLP
+
+- [Stanford Question Answering Dataset (SQuAD)](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/nlp/StanfordQuestionAnsweringDataset.html) - A reading comprehension dataset with text from wikipedia articles
+- [Tatoeba English French Dataset](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/nlp/TatoebaEnglishFrenchDataset.html) - An english-french translation dataset from the Tatoeba Project
+
+### Tabular
+
+- [Airfoil Self-Noise](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/tabular/AirfoilRandomAccess.html) - A 6 feature dataset from NASA tests of airfoils
+- [Ames House Pricing](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/tabular/AmesRandomAccess.html) - A 80 feature dataset to predict house prices
+
+### Time Series
+
+- [Daily Delhi Climate](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/tabular/DailyDelhiClimate.html)
\ No newline at end of file
diff --git a/docs/development/add_dataset_to_djl.md b/docs/development/add_dataset_to_djl.md
new file mode 100644
index 000000000..2660806d2
--- /dev/null
+++ b/docs/development/add_dataset_to_djl.md
@@ -0,0 +1,66 @@
+# Add a new dataset to DJL basic datasets
+
+This document outlines the procedure to add new datasets to DJL.
+
+## Step 1: Prepare the folder structure
+
+1. Navigate to the `test/resources/mlrepo/dataset` folder and create a folder in it to store your dataset based on its category.
+   For example, `cv/ai/djl/basicdataset/mnist`.
+2. Create a version folder within your newly created dataset's folder (e.g `0.0.1`). The version should match your dataset version.
+
+### Step 2: Create a `metadata.json` file
+
+You need to create a `metadata.json` file for the repository to load the dataset. You can refer to the format in the `metadata.json` files for existing datasets to create your own.
+
+**Note:** You need to update the sha1 hash of each file in your `metadata.json` file. Use the following command to get the sha1Hash value:
+
+```shell
+$ shasum -a 1 <file_name>
+```
+
+### Step 3: Create a Dataset implementation
+
+Create a class that implements the dataset and loads it.
+For more details on creating datasets, see the [dataset creation guide](how_to_use_dataset.md).
+You should also look at examples of official DJL datasets such as [`AmesRandomAccess`](https://github.com/deepjavalibrary/djl/blob/master/basicdataset/src/main/java/ai/djl/basicdataset/tabular/AmesRandomAccess.java)
+or [`Cifar10`](https://github.com/deepjavalibrary/djl/blob/master/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/Cifar10.java).
+
+Then, add some tests for the dataset.
+For testing, you can use a local repository such as:
+
+```java
+Repository repository = Repository.newInstace("testRepository", Paths.get("/test/resources/mlrepo"));
+```
+
+### Step 4: Update the datasets list
+
+Add your dataset to the [list of built-in datasets](../dataset.md).
+
+### Step 5: Upload metadata
+
+The official DJL ML repository is located on an S3 bucket managed by the AWS DJL team.
+You have to add the metadata and any dataset files to the repository.
+
+For non-AWS team members, go ahead straight to Step 6 and open a pull request.
+Within the pull request, you can coordinate with an AWS member to add the necessary files.
+
+For AWS team members, run the following command to upload your model to the S3 bucket:
+
+```shell
+$ ./gradlew syncS3
+```
+
+The `metadata.json` in DJL is mainly a repository of metadata.
+Within the metadata is typically contains only links indicating where the actual data would be found.
+
+However, some datasets can be distributed by DJL depending on whether it makes it easier to use and the dataset permits redistribution.
+In that case, coordinate with an AWS team member in your pull request.
+
+### Step 6: Open a PR to add your ModelLoader and metadata files to the git repository
+
+**Note**: Avoid checking in binary files to git. Binary files should only be uploaded to the S3 bucket.
+
+If you are relying on an AWS team member, you should leave your code with the local test repositories.
+If you try to use the official repository before it contains your metadata, the tests will not pass the CI.
+
+Once an AWS team member adds your metadata, they will prompt you to update your PR to the official repository.
diff --git a/docs/development/add_model_to_model-zoo.md b/docs/development/add_model_to_model-zoo.md
index 4dfa517a3..201bef465 100644
--- a/docs/development/add_model_to_model-zoo.md
+++ b/docs/development/add_model_to_model-zoo.md
@@ -1,6 +1,6 @@
-# Add a new model to the model zoo
+# Add a new model to the DJL model zoo
 
-This document outlines the procedure to add new models into the model zoo.
+This document outlines the procedure to add new models into the DJL model zoo.
 
 ## Step 1: Prepare the model files
 
@@ -31,6 +31,7 @@ For example, `image_classification/ai/djl/resnet`.
 3. Copy model files into the version folder.
 
 ### Step 3: Create a `metadata.json` file
+
 You need to create a `metadata.json` file for the model zoo to load the model. You can refer to the format in the `metadata.json` files for existing models to create your own.
 
 For a model built as a DJL block, you must recreate the block before loading the parameters. As part of your `metadata.json` file, you should use the `arguments` property to specify the arguments required for the model loader to create another `Block` matching the one used to train the model.
@@ -51,7 +52,7 @@ Verify that your folder has the following files (see Step 1 for additional files
 
 The official DJL ML repository is located on an S3 bucket managed by the AWS DJL team.
 
-For non-team members, coordinate with a team member in your pull request to coordinate adding the necessary files.
+For non-team members, coordinate with a team member in your pull request to add the necessary files.
 
 For AWS team members, run the following command to upload your model to the S3 bucket:
 
diff --git a/docs/development/example_dataset.md b/docs/development/example_dataset.md
new file mode 100644
index 000000000..78ed77a8c
--- /dev/null
+++ b/docs/development/example_dataset.md
@@ -0,0 +1,143 @@
+## Example CSV Dataset
+
+If the provided Datasets don't meet your requirements, you can also easily extend our dataset to create your own customized dataset.
+
+Let's take CSVDataset, which can load a csv file, for example.
+
+### Step 1: Prerequisites
+For this example, we'll use [malicious_url_data.csv](https://github.com/incertum/cyber-matrix-ai/blob/master/Malicious-URL-Detection-Deep-Learning/data/url_data_mega_deep_learning.csv).
+
+The CSV file has the following format.
+
+| URL      | isMalicious |
+| ----------- | ----------- |
+| sample.url.good.com | 0 |
+| sample.url.bad.com | 1  |
+
+We'll also use the 3rd party [Apache Commons](https://commons.apache.org/) library to read the CSV file. To use the library, include the following dependency:
+
+```
+api group: 'org.apache.commons', name: 'commons-csv', version: '1.7'
+```
+
+### Step 2: Implementation
+In order to extend the dataset, the following dependencies are required:
+
+```
+api "ai.djl:api:0.17.0"
+api "ai.djl:basicdataset:0.17.0"
+```
+
+There are four parts we need to implement for CSVDataset.
+
+1. Constructor and Builder
+
+First, we need a private field that holds the CSVRecord list from the csv file.
+We create a constructor and pass the CSVRecord list from builder to the class field.
+For builder, we have all we need in `BaseBuilder` so we only need to include the two minimal methods as shown.
+In the *build()* method, we take advantage of CSVParser to get the record of each CSV file and put them in CSVRecord list.
+
+```java
+public class CSVDataset extends RandomAccessDataset {
+
+    private final List<CSVRecord> csvRecords;
+
+    private CSVDataset(Builder builder) {
+        super(builder);
+        csvRecords = builder.csvRecords;
+    }
+    ...
+    public static final class Builder extends BaseBuilder<Builder> {
+        List<CSVRecord> csvRecords;
+
+        @Override
+        protected Builder self() {
+            return this;
+        }
+
+        CSVDataset build() throws IOException {
+            String csvFilePath = "path/malicious_url_data.csv";
+            try (Reader reader = Files.newBufferedReader(Paths.get(csvFilePath));
+                 CSVParser csvParser =
+                    new CSVParser(
+                        reader,
+                        CSVFormat.DEFAULT
+                            .withHeader("url", "isMalicious")
+                            .withFirstRecordAsHeader()
+                            .withIgnoreHeaderCase()
+                            .withTrim())) {
+                 csvRecords = csvParser.getRecords();
+            }
+            return new CSVDataset(this);
+        }
+    }
+
+}
+```
+
+2. Prepare
+
+As mentioned, in this example we are taking advantage of CSVParser to prepare the data for us. To prepare
+the data on our own, we use the `prepare()` method. Normally here we would load or create any data
+for our dataset and then save it in one of the private fields previously created. This `prepare()` method
+is called everytime we call `getData()` so in every case we want to only load the data once, we use a
+boolean variable called `prepared` to check if it has previously been loaded or prepared.
+
+Since we don't have to prepare any data on our own for this example, we only have to override it.
+
+```java
+@Override
+public void prepare(Progress progress) {}
+```
+
+There are great [examples](https://github.com/deepjavalibrary/djl/blob/master/basicdataset/src/main/java/ai/djl/basicdataset/nlp/AmazonReview.java)
+in our [basicdataset](https://github.com/deepjavalibrary/djl/blob/master/basicdataset/src/main/java/ai/djl/basicdataset)
+folder that show use cases for `prepare()`.
+
+
+
+3. Getter
+
+The getter returns a Record object which contains encoded inputs and labels.
+Here, we use simple encoding to transform the url String to an int array and create a NDArray on top of it.
+The reason why we use NDList here is that you might have multiple inputs and labels in different tasks.
+
+```java
+@Override
+public Record get(NDManager manager, long index) {
+    // get a CSVRecord given an index
+    CSVRecord record = csvRecords.get(Math.toIntExact(index));
+    NDArray datum = manager.create(encode(record.get("url")));
+    NDArray label = manager.create(Float.parseFloat(record.get("isMalicious")));
+    return new Record(new NDList(datum), new NDList(label));
+}
+```
+
+4. Size
+
+The number of records available to be read in this Dataset.
+Here, we can directly use the size of the List<CSVRecord>.
+
+```java
+@Override
+public long availableSize() {
+    return csvRecords.size();
+}
+```
+
+Done!
+Now, you can use the CSVDataset with the following code snippet:
+
+```java
+CSVDataset dataset = new CSVDataset.Builder().setSampling(batchSize, false).build();
+for (Batch batch : dataset.getData(model.getNDManager())) {
+    // use head to get first NDArray
+    batch.getData().head();
+    batch.getLabels().head();
+    ...
+    // don't forget to close the batch in the end
+    batch.close();
+}
+```
+
+Full example code could be found in [CSVDataset.java](https://github.com/deepjavalibrary/djl/blob/master/docs/development/CSVDataset.java).
diff --git a/docs/development/how_to_use_dataset.md b/docs/development/how_to_use_dataset.md
index 206957b04..8f6c3abd7 100644
--- a/docs/development/how_to_use_dataset.md
+++ b/docs/development/how_to_use_dataset.md
@@ -1,227 +1,58 @@
 # Dataset Creation
 
-The Dataset in DJL represents both the raw data and loading process.
-RandomAccessDataset implements the Dataset interface and provides comprehensive data loading functionality.
-RandomAccessDataset is also a basic dataset that supports random or sequential access of data using indices.
-You can easily customize your own dataset by extending RandomAccessDataset.
+The [Dataset](https://javadoc.io/doc/ai.djl/api/latest/ai/djl/training/dataset/Dataset.html) in DJL represents both the raw data and the data loading process.
+For this reason, training in DJL usually requires that your data be implemented through using a dataset class.
+You can choose to use one of the well-known datasets we have [built in](../dataset.md).
+Or, you can create a custom dataset.
 
-We provide several well-known datasets that you can use in our [Basic Datasets module](https://javadoc.io/doc/ai.djl/basicdataset/latest/index.html).
+## Dataset Helpers
 
-We also provide several built-in datasets that you can easily wrap around existing NDArrays and images.
+There are a number of helpers provided by DJL to make it easy to create custom datasets.
+If a helper is available, it can make it easier to implement the dataset then building it from scratch:
 
-- ArrayDataset - a dataset that wraps your existing NDArrays as inputs and labels.
-- ImageFolder - a dataset that wraps your images under folders for image classification.
+### CV
 
-If none of the provided datasets meet your requirements, you can also easily customize you own dataset by extending
-the RandomAccessDataset.
+- [ImageDataset](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/ImageDataset.html) - A abstract dataset to create a dataset where the input is an image such as image classification, object detection, and image segmentation
+- [ImageClassificationDataset](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/classification/ImageClassificationDataset.html) - An abstract dataset for image classification
+- [AbstractImageFolder](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/classification/AbstractImageFolder.html) - An abstract dataset for loading images in a folder structure. Usually you want the ImageFolderDataset.
+- [ImageFolder](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/classification/ImageFolder.html) - A dataset for loading image folders stored in a folder structure
+- [ObjectDetectionDataset](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/cv/ObjectDetectionDataset.html) - An abstract dataset for object detection
 
-## How to use the ArrayDataset
+### NLP
 
-The following code illustrates an implementation of [ArrayDataset](https://javadoc.io/doc/ai.djl/api/latest/ai/djl/training/dataset/ArrayDataset.html).
-The ArrayDataset is recommended only if your dataset is small enough to fit in memory.
+- [TextDataset](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/nlp/TextDataset.html) - An abstract dataset for NLP where either the input or labels are text-based.
+- [TextData](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/utils/TextData.html) - A utility for managing the text within a dataset
 
-```java
-// given you have data1, data2 and label1, label2, label3
-ArrayDataset dataset = new ArrayDataset.Builder()
-                              .setData(data1, data2)
-                              .optLabels(label1, label2, label3)
-                              .setSampling(20, false)
-                              .build();
+### Other
 
-```
+- [CsvDataset](https://javadoc.io/doc/ai.djl/basicdataset/latest/ai/djl/basicdataset/tabular/CsvDataset.html) - An dataset for loading data from a .csv file
 
-When you get the `Batch` from `trainer.iterateDataset(dataset)` or `dataset.getData(manager)`,
-you can use ``batch.getData()`` to get a NDList. The size of this NDList is the amount of data NDArrays 
-entered in `setData()`, in this case the size is 2. You can then use `NDList.get(0)` to get your first 
-data and `NDList.get(1)` to get your second data.
-Similarly, you can use `batch.getLabels()` to get a NDList with size 3.
+## Custom Datasets
 
-## How to use the ImageFolder dataset
+If none of the provided datasets meet your requirements, you can also easily customize you own dataset in a custom class.
+While technically the dataset must only implement [`Dataset`](https://javadoc.io/doc/ai.djl/api/latest/ai/djl/training/dataset/Dataset.html), it is best to instead extend [`RandomAccessDataset`](https://javadoc.io/doc/ai.djl/api/latest/ai/djl/training/dataset/RandomAccessDataset.html).
+It manages data randomization and provides comprehensive data loading functionality.
 
-This section outlines the procedure to use the ImageFolder dataset.
+The `RandomAccessDataset` is based on making your data records into a list where each record has an index.
+Then, it only needs to know how many records there are and how to load each record giving its index.
 
-The ImageFolder dataset is recommended only if you want to iterate through your images with their single label. For example, when training classic image classification.
+As part of implementing the dataset, there are two methods that must be defined:
 
-### Step 1: Prepare Images
-Arrange your image folder structure as follows:
+- `Record get(NDManager manager, long index)` - Returns the record (both input data and output label) for a particular index
+- `long availableSize()` - Returns the number of records in the dataset
 
-```
-dataset_root/shoes/Aerobic Shoes1.png
-dataset_root/shoes/Aerobic Shose2.png
-...
-dataset_root/boots/Black Boots.png
-dataset_root/boots/White Boots.png
-...
-dataset_root/pumps/Red Pumps
-dataset_root/pumps/Pink Pumps
-...
-```
-
-The dataset will take the folder name e.g. `boots`, `pumps`, `shoes` in sorted order as your labels
-**Note:** Nested folder structures are not currently supported
-
-### Step 2: Use the Dataset
-Add the following code snippet to your project to use the ImageFolder dataset.
+In addition, the dataset should also have a nested builder class to contain details on how to load the dataset.
+The builder would extend `RandomAccessDataset.BaseBuilder`.
+This provides an avenue to modify how RandomAccessDataset loads the data.
+You can also add your own options into the builder.
+For an example of how this would look like, see [`ImageFolder.Builder`](https://github.com/deepjavalibrary/djl/blob/master/basicdataset/src/main/java/ai/djl/basicdataset/cv/classification/ImageFolder.java).
 
-```java
-// set the image folder path
-Repository repository = Repository.newInstance("folder", Paths.get("/imagefolder");
-ImageFolder dataset =
-    new ImageFolder.Builder()
-        .setRepository(repository)
-        .addTransform(new Resize(100, 100))
-        .addTransform(new ToTensor())
-        .setSampling(batchSize, true)
-        .build();
-// call prepare before using
-dataset.prepare();
-
-// to get the synset
-List<String> synset = dataset.getSynset();
-```
+You can also view this example of creating a [new CSV dataset](example_dataset.md).
 
-Typically, you would add pre-processing pipelines like Resize(), in order to batchify the dataset, and ToTensor(), which converts the image NDArray to Tensor NDArray.
+Many of the abstract dataset helpers above also extend `RandomAccessDataset`.
+When using them, most of the same information applies.
+You may be asked to implement slightly different methods depending on the particular extended class.
+You will also want to extend that classes `BaseBuilder` instead of the one found in `RandomAccessDataset` to get the additional data loading options from the helper.
 
-## How to create a custom dataset
-
-If the provided Datasets don't meet your requirements, you can also easily extend our dataset to create your own customized dataset.
-
-Let's take CSVDataset, which can load a csv file, for example.
-
-### Step 1: Prerequisites
-For this example, we'll use [malicious_url_data.csv](https://github.com/incertum/cyber-matrix-ai/blob/master/Malicious-URL-Detection-Deep-Learning/data/url_data_mega_deep_learning.csv).
-
-The CSV file has the following format.
-
-| URL      | isMalicious |
-| ----------- | ----------- |
-| sample.url.good.com | 0 |
-| sample.url.bad.com | 1  |
-
-We'll also use the 3rd party [Apache Commons](https://commons.apache.org/) library to read the CSV file. To use the library, include the following dependency:
-
-```
-api group: 'org.apache.commons', name: 'commons-csv', version: '1.7'
-```
-
-### Step 2: Implementation
-In order to extend the dataset, the following dependencies are required:
-
-```
-api "ai.djl:api:0.17.0"
-api "ai.djl:basicdataset:0.17.0"
-```
-
-There are four parts we need to implement for CSVDataset.
-
-1. Constructor and Builder
-
-First, we need a private field that holds the CSVRecord list from the csv file.
-We create a constructor and pass the CSVRecord list from builder to the class field.
-For builder, we have all we need in `BaseBuilder` so we only need to include the two minimal methods as shown.
-In the *build()* method, we take advantage of CSVParser to get the record of each CSV file and put them in CSVRecord list.
-
-```java
-public class CSVDataset extends RandomAccessDataset {
-
-    private final List<CSVRecord> csvRecords;
-
-    private CSVDataset(Builder builder) {
-        super(builder);
-        csvRecords = builder.csvRecords;
-    }
-    ...
-    public static final class Builder extends BaseBuilder<Builder> {
-        List<CSVRecord> csvRecords;
-
-        @Override
-        protected Builder self() {
-            return this;
-        }
-
-        CSVDataset build() throws IOException {
-            String csvFilePath = "path/malicious_url_data.csv";
-            try (Reader reader = Files.newBufferedReader(Paths.get(csvFilePath));
-                 CSVParser csvParser =
-                    new CSVParser(
-                        reader,
-                        CSVFormat.DEFAULT
-                            .withHeader("url", "isMalicious")
-                            .withFirstRecordAsHeader()
-                            .withIgnoreHeaderCase()
-                            .withTrim())) {
-                 csvRecords = csvParser.getRecords();
-            }
-            return new CSVDataset(this);
-        }
-    }
-
-}
-```
-
-2. Prepare
-
-As mentioned, in this example we are taking advantage of CSVParser to prepare the data for us. To prepare
-the data on our own, we use the `prepare()` method. Normally here we would load or create any data 
-for our dataset and then save it in one of the private fields previously created. This `prepare()` method
-is called everytime we call `getData()` so in every case we want to only load the data once, we use a 
-boolean variable called `prepared` to check if it has previously been loaded or prepared.
-
-Since we don't have to prepare any data on our own for this example, we only have to override it.
-
-```java
-@Override
-public void prepare(Progress progress) {}
-```
-
-There are great [examples](https://github.com/deepjavalibrary/djl/blob/master/basicdataset/src/main/java/ai/djl/basicdataset/nlp/AmazonReview.java)
-in our [basicdataset](https://github.com/deepjavalibrary/djl/blob/master/basicdataset/src/main/java/ai/djl/basicdataset)
-folder that show use cases for `prepare()`.
-
-
-
-3. Getter
-
-The getter returns a Record object which contains encoded inputs and labels.
-Here, we use simple encoding to transform the url String to an int array and create a NDArray on top of it.
-The reason why we use NDList here is that you might have multiple inputs and labels in different tasks.
-
-```java
-@Override
-public Record get(NDManager manager, long index) {
-    // get a CSVRecord given an index
-    CSVRecord record = csvRecords.get(Math.toIntExact(index));
-    NDArray datum = manager.create(encode(record.get("url")));
-    NDArray label = manager.create(Float.parseFloat(record.get("isMalicious")));
-    return new Record(new NDList(datum), new NDList(label));
-}
-```
-
-4. Size
-
-The number of records available to be read in this Dataset.
-Here, we can directly use the size of the List<CSVRecord>.
-
-```java
-@Override
-public long availableSize() {
-    return csvRecords.size();
-}
-```
-
-Done!
-Now, you can use the CSVDataset with the following code snippet:
-
-```java
-CSVDataset dataset = new CSVDataset.Builder().setSampling(batchSize, false).build();
-for (Batch batch : dataset.getData(model.getNDManager())) {
-    // use head to get first NDArray
-    batch.getData().head();
-    batch.getLabels().head();
-    ...
-    // don't forget to close the batch in the end
-    batch.close();
-}
-```
-
-Full example code could be found in [CSVDataset.java](https://github.com/deepjavalibrary/djl/blob/master/docs/development/CSVDataset.java).
+If you create a new dataset for public dataset, consider contributing that dataset back to DJL for others to use.
+You can follow [these instructions](add_dataset_to_djl.md) for adding it.
\ No newline at end of file
diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml
index 041ae3d3e..026ccb953 100644
--- a/docs/mkdocs.yml
+++ b/docs/mkdocs.yml
@@ -54,6 +54,7 @@ nav:
           - 'docs/development/troubleshooting.md'
           - 'docs/development/dependency_management.md'
           - 'docs/development/add_model_to_model-zoo.md'
+          - 'docs/development/add_dataset_to_djl.md'
       - 'docs/roadmap.md'
       - 'docs/faq.md'
   - Tutorials:
@@ -66,7 +67,7 @@ nav:
       - 'jupyter/transfer_learning_on_cifar10.ipynb'
       - Load your own BERT:
           - BERT with MXNet: 'jupyter/mxnet/load_your_own_mxnet_bert.ipynb'
-          - BERT with PyTorch: 'jupyter/pytorch/load_your_own_pytorch.ipynb'
+          - BERT with PyTorch: 'jupyter/pytorch/load_your_own_pytorch_bert.ipynb'
   - Guides:
       - 'docs/engine.md'
       - Models:
diff --git a/engines/ml/xgboost/build.gradle b/engines/ml/xgboost/build.gradle
index 19affcc3a..9d31691b7 100644
--- a/engines/ml/xgboost/build.gradle
+++ b/engines/ml/xgboost/build.gradle
@@ -1,6 +1,8 @@
 import groovy.xml.QName
 
 group "ai.djl.ml.xgboost"
+boolean isGpu = project.hasProperty("gpu")
+def XGB_FLAVOR = isGpu ? "-gpu" : ""
 
 configurations {
     exclusion
@@ -9,7 +11,7 @@ configurations {
 dependencies {
     api project(":api")
     api "commons-logging:commons-logging:${commons_logging_version}"
-    api("ml.dmlc:xgboost4j_2.12:${xgboost_version}") {
+    api("ml.dmlc:xgboost4j${XGB_FLAVOR}_2.12:${xgboost_version}") {
         // get rid of the unused XGBoost Dependencies
         exclude group: "org.apache.hadoop", module: "hadoop-hdfs"
         exclude group: "org.apache.hadoop", module: "hadoop-common"
@@ -23,9 +25,13 @@ dependencies {
         exclude group: "org.scala-lang", module: "scala-reflect"
         exclude group: "org.scala-lang", module: "scala-library"
     }
+    if (isGpu) {
+        api "ai.rapids:cudf:${rapis_version}:cuda11"
+    }
 
     exclusion project(":api")
     exclusion "commons-logging:commons-logging:${commons_logging_version}"
+    exclusion "ai.rapids:cudf:${rapis_version}:cuda11"
     testImplementation(project(":testing"))
     testImplementation("org.testng:testng:${testng_version}") {
         exclude group: "junit", module: "junit"
@@ -43,6 +49,7 @@ jar {
                         "ml/dmlc/xgboost4j/java/NativeLibLoader*",
                         "ml/dmlc/xgboost4j/java/XGBoost*",
                         "ml/dmlc/xgboost4j/java/util/*",
+                        "ml/dmlc/xgboost4j/gpu/java/*",
                         "ml/dmlc/xgboost4j/LabeledPoint.*",
                         "xgboost4j-version.properties"
             }
@@ -53,6 +60,7 @@ jar {
 publishing {
     publications {
         maven(MavenPublication) {
+            artifactId "${project.name}${XGB_FLAVOR}"
             pom {
                 name = "DJL Engine Adapter for XGBoost"
                 description = "Deep Java Library (DJL) Engine Adapter for XGBoost"
diff --git a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbEngine.java b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbEngine.java
index 9aa03932a..e4a4ebc70 100644
--- a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbEngine.java
+++ b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbEngine.java
@@ -15,9 +15,13 @@ package ai.djl.ml.xgboost;
 import ai.djl.Device;
 import ai.djl.Model;
 import ai.djl.engine.Engine;
+import ai.djl.engine.StandardCapabilities;
 import ai.djl.ndarray.NDManager;
 import ai.djl.nn.SymbolBlock;
 import ai.djl.training.GradientCollector;
+import java.io.IOException;
+import java.io.InputStream;
+import java.util.Properties;
 import ml.dmlc.xgboost4j.java.JniUtils;
 
 /**
@@ -71,12 +75,27 @@ public final class XgbEngine extends Engine {
     /** {@inheritDoc} */
     @Override
     public String getVersion() {
-        return "1.3.1";
+        try (InputStream is =
+                XgbEngine.class.getResourceAsStream("/xgboost4j-version.properties")) {
+            Properties prop = new Properties();
+            prop.load(is);
+            return prop.getProperty("version");
+        } catch (IOException e) {
+            throw new AssertionError("Failed to load xgboost4j-version.properties", e);
+        }
     }
 
     /** {@inheritDoc} */
     @Override
     public boolean hasCapability(String capability) {
+        if (StandardCapabilities.CUDA.equals(capability)) {
+            try {
+                Class.forName("ml.dmlc.xgboost4j.gpu.java.CudfColumn");
+                return true;
+            } catch (ClassNotFoundException ignore) {
+                return false;
+            }
+        }
         return false;
     }
 
diff --git a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDManager.java b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDManager.java
index 87b6af06a..00b5133e2 100644
--- a/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDManager.java
+++ b/engines/ml/xgboost/src/main/java/ai/djl/ml/xgboost/XgbNDManager.java
@@ -20,11 +20,14 @@ import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.ndarray.types.SparseFormat;
+import ai.djl.util.JsonUtils;
+import com.google.gson.JsonArray;
 import java.nio.Buffer;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.nio.FloatBuffer;
 import java.util.Arrays;
+import ml.dmlc.xgboost4j.java.ColumnBatch;
 import ml.dmlc.xgboost4j.java.JniUtils;
 
 /** {@code XgbNDManager} is the XGBoost implementation of {@link NDManager}. */
@@ -80,6 +83,30 @@ public class XgbNDManager extends BaseNDManager {
         return Engine.getEngine(XgbEngine.ENGINE_NAME);
     }
 
+    /**
+     * Creates {@link XgbNDArray} from column array interface.
+     *
+     * @param columnBatch – the XGBoost ColumnBatch to provide the cuda array interface of feature
+     *     columns
+     * @param missing – missing value
+     * @param nthread – threads number
+     * @return a new instance of {@link NDArray}
+     */
+    public NDArray create(ColumnBatch columnBatch, float missing, int nthread) {
+        columnBatch.getFeatureArrayInterface();
+        String json = columnBatch.getFeatureArrayInterface();
+        JsonArray array = JsonUtils.GSON.fromJson(json, JsonArray.class);
+        JsonArray shapeJson = array.get(0).getAsJsonObject().get("shape").getAsJsonArray();
+        long[] shapes = new long[shapeJson.size()];
+        for (int i = 0; i < shapes.length; ++i) {
+            shapes[i] = shapeJson.get(i).getAsLong();
+        }
+
+        Shape shape = new Shape(shapes);
+        long handle = JniUtils.createDMatrix(columnBatch, missing, nthread);
+        return new XgbNDArray(this, alternativeManager, handle, shape, SparseFormat.DENSE);
+    }
+
     /** {@inheritDoc} */
     @Override
     public NDArray create(Buffer data, Shape shape, DataType dataType) {
diff --git a/engines/ml/xgboost/src/main/java/ml/dmlc/xgboost4j/java/JniUtils.java b/engines/ml/xgboost/src/main/java/ml/dmlc/xgboost4j/java/JniUtils.java
index ce63ad8bf..e660150ac 100644
--- a/engines/ml/xgboost/src/main/java/ml/dmlc/xgboost4j/java/JniUtils.java
+++ b/engines/ml/xgboost/src/main/java/ml/dmlc/xgboost4j/java/JniUtils.java
@@ -52,6 +52,19 @@ public final class JniUtils {
         return handles[0];
     }
 
+    public static long createDMatrix(ColumnBatch columnBatch, float missing, int nthread) {
+        long[] handles = new long[1];
+        String json = columnBatch.getFeatureArrayInterface();
+        if (json == null || json.isEmpty()) {
+            throw new IllegalArgumentException(
+                    "Expecting non-empty feature columns' array interface");
+        }
+        checkCall(
+                XGBoostJNI.XGDMatrixCreateFromArrayInterfaceColumns(
+                        json, missing, nthread, handles));
+        return handles[0];
+    }
+
     public static long createDMatrixCSR(long[] indptr, int[] indices, float[] array) {
         long[] handles = new long[1];
         checkCall(XGBoostJNI.XGDMatrixCreateFromCSREx(indptr, indices, array, 0, handles));
diff --git a/engines/mxnet/mxnet-engine/README.md b/engines/mxnet/mxnet-engine/README.md
index 50d1e4df1..5713b8712 100644
--- a/engines/mxnet/mxnet-engine/README.md
+++ b/engines/mxnet/mxnet-engine/README.md
@@ -41,7 +41,10 @@ You can pull the MXNet engine from the central Maven repository by including the
 By default, DJL will download the Apache MXNet native libraries into [cache folder](../../../docs/development/cache_management.md) the first time you run DJL.
 It will automatically determine the appropriate jars for your system based on the platform and GPU support.
 
-You can choose a native library based on your platform if you don't have network access at runtime.
+If you do not want to rely on the download because you don't have network access at runtime or for other reasons, there are additional options.
+The easiest option is to add a DJL native library package to your project dependencies.
+The available packages for your platform can be found below.
+Finally, you can also specify the path to a valid MXNet build using the `MXNET_LIBRARY_PATH` environment variable.
 
 ### macOS
 
@@ -116,15 +119,14 @@ Apache MXNet requires Visual C++ Redistributable Packages. If you encounter an U
 DJL on Windows, please download and install
 [Visual C++ 2019 Redistributable Packages](https://support.microsoft.com/en-us/help/2977003/the-latest-supported-visual-c-downloads) and reboot.
 
-For the Windows platform, you can use CPU package. MXNet windows GPU native
-library size are large, we no longer provide GPU package, instead you have to
-use the Automatic package.
+For the windows platform, we support both CPU and GPU.
+The CPU can be found using either the automatic runtime detection or through adding the CPU jar to your dependencies.
+However, due to the size of the windows GPU native library, we do not offer GPU support through a dependency jar.
+You can still access GPU on windows by using the [automatic runtime download](#Installation).
 
 #### Windows GPU
 
-- ai.djl.mxnet:mxnet-native-auto:1.8.0
-
-    This package supports CUDA 11.0 and CUDA 10.2 for Windows.
+This package supports CUDA 11.0 and CUDA 10.2 for Windows.
 
 ### Windows CPU
 
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
index 66fe02d5d..fc68b4b73 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArray.java
@@ -315,7 +315,9 @@ public class MxNDArray extends NativeResource<Pointer> implements LazyNDArray {
     /** {@inheritDoc} */
     @Override
     public NDArray take(NDArray index) {
-        throw new UnsupportedOperationException("Not implemented yet.");
+        MxOpParams params = new MxOpParams();
+        params.add("mode", "wrap");
+        return manager.invoke("take", new NDList(this.flatten(), index), params).singletonOrThrow();
     }
 
     /** {@inheritDoc} */
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayEx.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayEx.java
index aa5d58948..7eaac6ab8 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayEx.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayEx.java
@@ -984,8 +984,8 @@ class MxNDArrayEx implements NDArrayEx {
 
     /** {@inheritDoc} */
     @Override
-    public NDArrayIndexer getIndexer() {
-        return new MxNDArrayIndexer(array.getManager());
+    public NDArrayIndexer getIndexer(NDManager manager) {
+        return new MxNDArrayIndexer((MxNDManager) manager);
     }
 
     ////////////////////////////////////////
diff --git a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayIndexer.java b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayIndexer.java
index 668d70d30..e5e1f09dc 100644
--- a/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayIndexer.java
+++ b/engines/mxnet/mxnet-engine/src/main/java/ai/djl/mxnet/engine/MxNDArrayIndexer.java
@@ -37,8 +37,7 @@ public class MxNDArrayIndexer extends NDArrayIndexer {
         params.addParam("axis", fullPick.getAxis());
         params.addParam("keepdims", true);
         params.add("mode", "wrap");
-        return array.getManager()
-                .invoke("pick", new NDList(array, fullPick.getIndices()), params)
+        return manager.invoke("pick", new NDList(array, fullPick.getIndices()), params)
                 .singletonOrThrow();
     }
 
@@ -51,7 +50,7 @@ public class MxNDArrayIndexer extends NDArrayIndexer {
         params.addTupleParam("end", fullSlice.getMax());
         params.addTupleParam("step", fullSlice.getStep());
 
-        NDArray result = ((MxNDManager) array.getManager()).invoke("_npi_slice", array, params);
+        NDArray result = manager.invoke("_npi_slice", array, params);
         int[] toSqueeze = fullSlice.getToSqueeze();
         if (toSqueeze.length > 0) {
             NDArray oldResult = result;
@@ -83,12 +82,11 @@ public class MxNDArrayIndexer extends NDArrayIndexer {
         prepareValue.add(prepareValue.peek().reshape(targetShape));
         prepareValue.add(prepareValue.peek().broadcast(fullSlice.getShape()));
 
-        array.getManager()
-                .invoke(
-                        "_npi_slice_assign",
-                        new NDArray[] {array, prepareValue.peek()},
-                        new NDArray[] {array},
-                        params);
+        manager.invoke(
+                "_npi_slice_assign",
+                new NDArray[] {array, prepareValue.peek()},
+                new NDArray[] {array},
+                params);
         for (NDArray toClean : prepareValue) {
             if (toClean != value) {
                 toClean.close();
@@ -105,11 +103,7 @@ public class MxNDArrayIndexer extends NDArrayIndexer {
         params.addTupleParam("end", fullSlice.getMax());
         params.addTupleParam("step", fullSlice.getStep());
         params.addParam("scalar", value);
-        array.getManager()
-                .invoke(
-                        "_npi_slice_assign_scalar",
-                        new NDArray[] {array},
-                        new NDArray[] {array},
-                        params);
+        manager.invoke(
+                "_npi_slice_assign_scalar", new NDArray[] {array}, new NDArray[] {array}, params);
     }
 }
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
index a8944cd92..96d30d87c 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArray.java
@@ -235,8 +235,8 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
 
     /** {@inheritDoc} */
     @Override
-    public NDArray get(long... indices) {
-        return JniUtils.getItem(this, indices);
+    public NDArray get(NDManager manager, long... indices) {
+        return JniUtils.getItem(this, indices, (PtNDManager) manager);
     }
 
     /** {@inheritDoc} */
@@ -274,8 +274,8 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
     /** {@inheritDoc} */
     @Override
     public void tempAttach(NDManager manager) {
-        detach();
         NDManager original = this.manager;
+        detach();
         this.manager = (PtNDManager) manager;
         manager.tempAttachInternal(original, getUid(), this);
     }
@@ -301,7 +301,7 @@ public class PtNDArray extends NativeResource<Long> implements NDArray {
             // Result is flattened since shape is undetermined
             return JniUtils.booleanMask(this, manager.from(index));
         } else if (indexShape.equals(getShape().slice(axis))) {
-            // index will be broadcasted by default
+            // index will be broadcast by default
             try (PtNDArray flattedResult = JniUtils.booleanMask(this, manager.from(index))) {
                 // Shape recovery
                 Shape remainder = getShape().slice(0, axis);
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayEx.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayEx.java
index ffa4a8da1..3f4d5ffd0 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayEx.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayEx.java
@@ -624,8 +624,8 @@ public class PtNDArrayEx implements NDArrayEx {
 
     /** {@inheritDoc} */
     @Override
-    public NDArrayIndexer getIndexer() {
-        return new PtNDArrayIndexer(array.getManager());
+    public NDArrayIndexer getIndexer(NDManager manager) {
+        return new PtNDArrayIndexer((PtNDManager) manager);
     }
 
     /** {@inheritDoc} */
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
index bd9f465a5..f86575618 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/engine/PtNDArrayIndexer.java
@@ -43,7 +43,7 @@ public class PtNDArrayIndexer extends NDArrayIndexer {
         long[] min = fullSlice.getMin();
         long[] max = fullSlice.getMax();
         long[] step = fullSlice.getStep();
-        try (PtNDArray res = JniUtils.index(manager.from(array), min, max, step)) {
+        try (PtNDArray res = JniUtils.index(manager.from(array), min, max, step, manager)) {
             return res.squeeze(fullSlice.getToSqueeze());
         }
     }
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 5d0f21b61..5a0cf2a82 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
@@ -324,9 +324,13 @@ public final class JniUtils {
     }
 
     public static PtNDArray index(
-            PtNDArray ndArray, long[] minIndices, long[] maxIndices, long[] stepIndices) {
+            PtNDArray ndArray,
+            long[] minIndices,
+            long[] maxIndices,
+            long[] stepIndices,
+            PtNDManager manager) {
         return new PtNDArray(
-                ndArray.getManager(),
+                manager,
                 PyTorchLibrary.LIB.torchIndex(
                         ndArray.getHandle(), minIndices, maxIndices, stepIndices));
     }
@@ -413,18 +417,16 @@ public final class JniUtils {
                 ndArray.getHandle(), value.getHandle(), indicesNd.getHandle());
     }
 
-    public static PtNDArray getItem(PtNDArray ndArray, long[] indices) {
+    public static PtNDArray getItem(PtNDArray ndArray, long[] indices, PtNDManager manager) {
         // use a specialized API here
         // due to significant performance gain
         // for commonly used data loading call
         if (indices.length == 1) {
             return new PtNDArray(
-                    ndArray.getManager(),
-                    PyTorchLibrary.LIB.torchGetItem(ndArray.getHandle(), indices[0]));
+                    manager, PyTorchLibrary.LIB.torchGetItem(ndArray.getHandle(), indices[0]));
         }
         return new PtNDArray(
-                ndArray.getManager(),
-                PyTorchLibrary.LIB.torchGetItem(ndArray.getHandle(), indices));
+                manager, PyTorchLibrary.LIB.torchGetItem(ndArray.getHandle(), indices));
     }
 
     public static PtNDArray clone(PtNDArray ndArray) {
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
index 421cd1fe7..563fefd9f 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArray.java
@@ -229,8 +229,8 @@ public class TfNDArray extends NativeResource<TFE_TensorHandle> implements NDArr
     /** {@inheritDoc} */
     @Override
     public void tempAttach(NDManager manager) {
-        detach();
         NDManager original = this.manager;
+        detach();
         this.manager = (TfNDManager) manager;
         manager.tempAttachInternal(original, getUid(), this);
     }
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayEx.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayEx.java
index 76ab99bf6..52406eaae 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayEx.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayEx.java
@@ -15,6 +15,7 @@ package ai.djl.tensorflow.engine;
 
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
 import ai.djl.ndarray.NDUtils;
 import ai.djl.ndarray.index.NDArrayIndexer;
 import ai.djl.ndarray.internal.NDArrayEx;
@@ -555,8 +556,8 @@ public class TfNDArrayEx implements NDArrayEx {
 
     /** {@inheritDoc} */
     @Override
-    public NDArrayIndexer getIndexer() {
-        return new TfNDArrayIndexer(array.getManager());
+    public NDArrayIndexer getIndexer(NDManager manager) {
+        return new TfNDArrayIndexer((TfNDManager) manager);
     }
 
     /** {@inheritDoc} */
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayIndexer.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayIndexer.java
index 4b2cdc744..73b3a4c5d 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayIndexer.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfNDArrayIndexer.java
@@ -36,14 +36,12 @@ public class TfNDArrayIndexer extends NDArrayIndexer {
     @Override
     public NDArray get(NDArray array, NDIndexFullSlice fullSlice) {
         array = manager.from(array);
-        TfNDManager tfManager = (TfNDManager) array.getManager();
         int[] toSqueeze = fullSlice.getToSqueeze();
-        try (NDArray begin = tfManager.create(fullSlice.getMin());
-                NDArray end = tfManager.create(fullSlice.getMax());
-                NDArray step = tfManager.create(fullSlice.getStep())) {
+        try (NDArray begin = manager.create(fullSlice.getMin());
+                NDArray end = manager.create(fullSlice.getMax());
+                NDArray step = manager.create(fullSlice.getStep())) {
             NDArray result =
-                    tfManager
-                            .opExecutor("StridedSlice")
+                    manager.opExecutor("StridedSlice")
                             .addInput(array)
                             .addInput(begin)
                             .addInput(end)
diff --git a/examples/src/main/java/ai/djl/examples/training/transferlearning/TrainAmazonReviewRanking.java b/examples/src/main/java/ai/djl/examples/training/transferlearning/TrainAmazonReviewRanking.java
index c2ef09933..462292c69 100644
--- a/examples/src/main/java/ai/djl/examples/training/transferlearning/TrainAmazonReviewRanking.java
+++ b/examples/src/main/java/ai/djl/examples/training/transferlearning/TrainAmazonReviewRanking.java
@@ -17,6 +17,8 @@ import ai.djl.Model;
 import ai.djl.ModelException;
 import ai.djl.basicdataset.tabular.CsvDataset;
 import ai.djl.basicdataset.utils.DynamicBuffer;
+import ai.djl.basicdataset.utils.Feature;
+import ai.djl.basicdataset.utils.Featurizer;
 import ai.djl.engine.Engine;
 import ai.djl.examples.training.util.Arguments;
 import ai.djl.inference.Predictor;
@@ -126,11 +128,9 @@ public final class TrainAmazonReviewRanking {
                 .optCsvUrl(amazonReview)
                 .setCsvFormat(CSVFormat.TDF.builder().setQuote(null).setHeader().build())
                 .setSampling(arguments.getBatchSize(), true)
-                .addFeature(
-                        new CsvDataset.Feature(
-                                "review_body", new BertFeaturizer(tokenizer, maxLength)))
+                .addFeature(new Feature("review_body", new BertFeaturizer(tokenizer, maxLength)))
                 .addLabel(
-                        new CsvDataset.Feature(
+                        new Feature(
                                 "star_rating",
                                 (buf, data) -> buf.put(Float.parseFloat(data) - 1.0f)))
                 .optDataBatchifier(
@@ -211,7 +211,7 @@ public final class TrainAmazonReviewRanking {
                 .addTrainingListeners(listener);
     }
 
-    private static final class BertFeaturizer implements CsvDataset.Featurizer {
+    private static final class BertFeaturizer implements Featurizer {
 
         private final BertFullTokenizer tokenizer;
         private final int maxLength;
diff --git a/extensions/audio/build.gradle b/extensions/audio/build.gradle
new file mode 100644
index 000000000..8b6d60530
--- /dev/null
+++ b/extensions/audio/build.gradle
@@ -0,0 +1,46 @@
+group 'ai.djl.audio'
+
+dependencies {
+    api project(":api")
+    api project(":basicdataset")
+    api ("org.bytedeco:javacv:1.5.7") {
+        exclude group: "org.bytedeco", module: "opencv"
+        exclude group: "org.bytedeco", module: "openblas"
+        exclude group: "org.bytedeco", module: "flycapture"
+        exclude group: "org.bytedeco", module: "libdc1394"
+        exclude group: "org.bytedeco", module: "libfreenect"
+        exclude group: "org.bytedeco", module: "libfreenect2"
+        exclude group: "org.bytedeco", module: "librealsense"
+        exclude group: "org.bytedeco", module: "librealsense2"
+        exclude group: "org.bytedeco", module: "videoinput"
+        exclude group: "org.bytedeco", module: "artoolkitplus"
+        exclude group: "org.bytedeco", module: "flandmark"
+        exclude group: "org.bytedeco", module: "leptonica"
+        exclude group: "org.bytedeco", module: "tesseract"
+        exclude group: "org.bytedeco", module: "tesseract"
+    }
+    api "com.github.wendykierp:JTransforms:3.1"
+
+    runtimeOnly "org.bytedeco:ffmpeg:5.0-1.5.7:macosx-x86_64"
+    runtimeOnly "org.bytedeco:ffmpeg:5.0-1.5.7:linux-x86_64"
+    runtimeOnly "org.bytedeco:ffmpeg:5.0-1.5.7:windows-x86_64"
+
+    testImplementation("org.testng:testng:${testng_version}") {
+        exclude group: "junit", module: "junit"
+    }
+    testImplementation project(":testing")
+    testRuntimeOnly "org.apache.logging.log4j:log4j-slf4j-impl:${log4j_slf4j_version}"
+    testRuntimeOnly project(":engines:pytorch:pytorch-engine")
+}
+
+publishing {
+    publications {
+        maven(MavenPublication) {
+            pom {
+                name = "DJL Audio processing extension"
+                description = "DJL Audio processing extension"
+                url = "http://www.djl.ai/extensions/${project.name}"
+            }
+        }
+    }
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/AudioUtils.java b/extensions/audio/src/main/java/ai/djl/audio/AudioUtils.java
new file mode 100644
index 000000000..e2a2bcb7d
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/AudioUtils.java
@@ -0,0 +1,32 @@
+/*
+ * 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.audio;
+
+import ai.djl.ndarray.NDArray;
+
+/** Useful utils for audio process. */
+public final class AudioUtils {
+
+    private AudioUtils() {}
+
+    /**
+     * Calculate root mean square energy in decibels.
+     *
+     * @param samples input signal, should be 1 dimension array
+     * @return root mean square energy
+     */
+    public static float rmsDb(NDArray samples) {
+        samples = samples.pow(2).mean().log10().mul(10);
+        return samples.toFloatArray()[0];
+    }
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/dataset/AudioData.java b/extensions/audio/src/main/java/ai/djl/audio/dataset/AudioData.java
new file mode 100644
index 000000000..bf8874e7d
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/dataset/AudioData.java
@@ -0,0 +1,223 @@
+/*
+ * 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.audio.dataset;
+
+import ai.djl.audio.processor.AudioNormalizer;
+import ai.djl.audio.processor.AudioProcessor;
+import ai.djl.audio.processor.LinearSpecgram;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+import ai.djl.repository.LocalRepository;
+import java.nio.Buffer;
+import java.nio.ShortBuffer;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.List;
+import org.bytedeco.javacv.FFmpegFrameGrabber;
+import org.bytedeco.javacv.Frame;
+import org.bytedeco.javacv.FrameGrabber;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * {@link AudioData} is a utility for managing audio data within a {@link
+ * ai.djl.training.dataset.Dataset}. It contains some basic information of an audio file and provide
+ * some method to preprocess the original audio file. Since storing all the data into the memory is
+ * impossible, this class will only store the path of the original audio data.
+ *
+ * <p>This class provides a list of {@link AudioProcessor} for user to featurize the data.
+ *
+ * <p>See {@link SpeechRecognitionDataset} for an example.
+ */
+public class AudioData {
+
+    private static final Logger logger = LoggerFactory.getLogger(LocalRepository.class);
+
+    private int sampleRate;
+    private int audioChannels;
+
+    private List<AudioProcessor> processorList;
+    private List<String> audioPaths;
+
+    /**
+     * Constructs a new {@link AudioData}.
+     *
+     * @param configuration the configuration for the {@link AudioData}.
+     */
+    public AudioData(Configuration configuration) {
+        this.sampleRate = configuration.sampleRate;
+        this.processorList = configuration.processorList;
+    }
+
+    /**
+     * Returns a good default {@link AudioData.Configuration} to use for the constructor with
+     * defaults.
+     *
+     * @return a good default {@link AudioData.Configuration} to use for the constructor with
+     *     defaults.
+     */
+    public static AudioData.Configuration getDefaultConfiguration() {
+        float targetDb = -20f;
+        float strideMs = 10f;
+        float windowsMs = 20f;
+        int sampleRate = 16000;
+        List<AudioProcessor> defaultProcessors =
+                Arrays.asList(
+                        new AudioNormalizer(targetDb),
+                        new LinearSpecgram(strideMs, windowsMs, sampleRate));
+        return new AudioData.Configuration()
+                .setProcessorList(defaultProcessors)
+                .setSampleRate(sampleRate);
+    }
+
+    /**
+     * This method is used for decoding the original audio data and converting it to a float array.
+     *
+     * @param path The path of the original audio data.
+     * @return A float array.
+     */
+    private float[] toFloat(String path) {
+        List<Float> list = new ArrayList<>();
+        float scale = (float) 1.0 / (float) (1 << (8 * 2) - 1);
+        try (FFmpegFrameGrabber audioGrabber = new FFmpegFrameGrabber(path)) {
+            audioGrabber.start();
+            audioChannels = audioGrabber.getAudioChannels();
+            audioGrabber.setSampleRate(sampleRate);
+            Frame frame;
+            while ((frame = audioGrabber.grabFrame()) != null) {
+                Buffer[] buffers = frame.samples;
+                ShortBuffer sb = (ShortBuffer) buffers[0];
+                for (int i = 0; i < sb.limit(); i++) {
+                    list.add(sb.get() * scale);
+                }
+            }
+        } catch (FrameGrabber.Exception e) {
+            logger.error(e.getMessage());
+        }
+        float[] floatArray = new float[list.size()];
+        int i = 0;
+        for (Float f : list) {
+            floatArray[i++] = (f != null ? f : Float.NaN);
+        }
+        return floatArray;
+    }
+
+    /**
+     * This method will use a list of {@link AudioProcessor} as featurizer to process the float
+     * array.
+     *
+     * @param manager the manager for converting the float array to {@link NDArray}.
+     * @param index The index of path of the original audio data.
+     * @return An {@link NDArray} that represent the processed audio data.
+     */
+    public NDArray getPreprocessedData(NDManager manager, int index) {
+        float[] floatArray = toFloat(audioPaths.get(index));
+        NDArray samples = manager.create(floatArray);
+        for (AudioProcessor processor : processorList) {
+            samples = processor.extractFeatures(manager, samples);
+        }
+        return samples;
+    }
+
+    /**
+     * Get the number of channels of an audio file.
+     *
+     * @return The number of channels of an audio file.
+     */
+    public int getAudioChannels() {
+        return audioChannels;
+    }
+
+    /**
+     * Get the sample rate used by {@link FFmpegFrameGrabber} when sampling the audio file.
+     *
+     * @return The sample rate used by {@link FFmpegFrameGrabber} when sampling the audio file.
+     */
+    public int getSampleRate() {
+        return sampleRate;
+    }
+
+    /**
+     * Set the path list of original audio data.
+     *
+     * @param audioPaths The path list of original audio data.
+     */
+    public void setAudioPaths(List<String> audioPaths) {
+        this.audioPaths = audioPaths;
+    }
+
+    /**
+     * Get the original audio path.
+     *
+     * @return The original audio path.
+     */
+    public List<String> getAudioPaths() {
+        return audioPaths;
+    }
+
+    /**
+     * Get the total number of audio data in the dataset.
+     *
+     * @return The total number of audio data in the dataset.
+     */
+    public int getTotalSize() {
+        return audioPaths.size();
+    }
+
+    /**
+     * The configuration for creating a {@link AudioData} value in a {@link
+     * ai.djl.audio.dataset.AudioData}.
+     */
+    public static final class Configuration {
+
+        private int sampleRate;
+
+        private List<AudioProcessor> processorList;
+
+        /**
+         * Set the list of processor which are used for extracting features from audio data.
+         *
+         * @param processorList The list of processor which are used for extracting features from
+         *     audio data.
+         * @return this configuration.
+         */
+        public Configuration setProcessorList(List<AudioProcessor> processorList) {
+            this.processorList = processorList;
+            return this;
+        }
+
+        /**
+         * Set the sampleRate for {@link FFmpegFrameGrabber} to use.
+         *
+         * @param sampleRate The sampleRate for {@link FFmpegFrameGrabber} to use.
+         * @return this configuration.
+         */
+        public Configuration setSampleRate(int sampleRate) {
+            this.sampleRate = sampleRate;
+            return this;
+        }
+
+        /**
+         * Updates this {@link AudioData.Configuration} with the non-null values from another
+         * configuration.
+         *
+         * @param other the other configuration to use to update this
+         * @return this configuration after updating
+         */
+        public AudioData.Configuration update(AudioData.Configuration other) {
+            processorList = other.processorList;
+            return this;
+        }
+    }
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/dataset/Librispeech.java b/extensions/audio/src/main/java/ai/djl/audio/dataset/Librispeech.java
new file mode 100644
index 000000000..9dd24aec0
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/dataset/Librispeech.java
@@ -0,0 +1,187 @@
+/*
+ * 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.audio.dataset;
+
+import ai.djl.Application;
+import ai.djl.basicdataset.BasicDatasets;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.repository.Artifact;
+import ai.djl.repository.MRL;
+import ai.djl.training.dataset.Dataset;
+import ai.djl.training.dataset.Record;
+import ai.djl.translate.TranslateException;
+import ai.djl.util.Progress;
+import java.io.BufferedReader;
+import java.io.File;
+import java.io.IOException;
+import java.nio.file.Files;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ * LibriSpeech is a corpus of approximately 1000 hours of 16kHz read English speech, prepared by
+ * Vassil Panayotov with the assistance of Daniel Povey. The data is derived from read audiobooks
+ * from the LibriVox project, and has been carefully segmented and aligned.
+ */
+public class Librispeech extends SpeechRecognitionDataset {
+
+    private static final String VERSION = "1.0";
+    private static final String ARTIFACT_ID = "librispeech";
+
+    /**
+     * Creates a new instance of {@link SpeechRecognitionDataset} with the given necessary
+     * configurations.
+     *
+     * @param builder a builder with the necessary configurations
+     */
+    public Librispeech(Builder builder) {
+        super(builder);
+        this.usage = builder.usage;
+        this.mrl = builder.getMrl();
+    }
+
+    /**
+     * Creates a builder to build a {@link Librispeech}.
+     *
+     * @return a new {@link Librispeech.Builder} object
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /**
+     * Prepares the dataset for use with tracked progress.
+     *
+     * @param progress the progress tracker
+     * @throws IOException for various exceptions depending on the dataset
+     */
+    @Override
+    public void prepare(Progress progress) throws IOException, TranslateException {
+        if (prepared) {
+            return;
+        }
+        Artifact artifact = mrl.getDefaultArtifact();
+        mrl.prepare(artifact, progress);
+        Artifact.Item item;
+        String subPath;
+        switch (usage) {
+            case TRAIN:
+                item = artifact.getFiles().get("train");
+                subPath = "LibriSpeech/test-clean-100";
+                break;
+            case TEST:
+                item = artifact.getFiles().get("test");
+                subPath = "LibriSpeech/test-clean";
+                break;
+            default:
+                throw new UnsupportedOperationException("Unsupported usage type.");
+        }
+        File mainDir = mrl.getRepository().getFile(item, subPath).toFile();
+        File[] subDirs = mainDir.listFiles();
+        if (subDirs == null) {
+            return;
+        }
+        List<String> lineArray = new ArrayList<>();
+        List<String> audioPaths = new ArrayList<>();
+        for (File subDir : subDirs) {
+            File[] subSubDirs = subDir.listFiles();
+            String subDirName = subDir.getName();
+            if (subSubDirs == null) {
+                return;
+            }
+            for (File subSubDir : subSubDirs) {
+                String subSubDirName = subSubDir.getName();
+                File transFile =
+                        new File(
+                                String.format(
+                                        "%s/%s-%s.trans.txt",
+                                        subSubDir.getAbsolutePath(), subDirName, subSubDirName));
+                try (BufferedReader reader = Files.newBufferedReader(transFile.toPath())) {
+                    String row;
+                    while ((row = reader.readLine()) != null) {
+                        if (row.contains(" ")) {
+                            String trans = row.substring(row.indexOf(' ') + 1);
+                            String label = row.substring(0, row.indexOf(' '));
+                            String audioIndex = label.split("-")[2];
+                            String audioPath =
+                                    String.format(
+                                            "%s/%s-%s-%s.flac",
+                                            subSubDir.getAbsolutePath(),
+                                            subDirName,
+                                            subSubDirName,
+                                            audioIndex);
+                            audioPaths.add(audioPath);
+                            lineArray.add(trans);
+                        }
+                    }
+                }
+            }
+        }
+        targetPreprocess(lineArray);
+        sourcePreprocess(audioPaths);
+        prepared = true;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Record get(NDManager manager, long index) {
+        NDList data = new NDList();
+        NDList labels = new NDList();
+        data.add(sourceAudioData.getPreprocessedData(manager, (int) index));
+        labels.add(targetTextData.getEmbedding(manager, index));
+        return new Record(data, labels);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected long availableSize() {
+        return sourceAudioData.getTotalSize();
+    }
+
+    /** A builder to construct a {@link Librispeech} . */
+    public static class Builder extends AudioBuilder<Librispeech.Builder> {
+
+        /** Constructs a new builder. */
+        public Builder() {
+            repository = BasicDatasets.REPOSITORY;
+            groupId = BasicDatasets.GROUP_ID;
+            artifactId = ARTIFACT_ID;
+            usage = Dataset.Usage.TRAIN;
+        }
+
+        /**
+         * Builds a new {@link Librispeech} object.
+         *
+         * @return the new {@link Librispeech} object
+         */
+        public Librispeech build() {
+            return new Librispeech(this);
+        }
+
+        /**
+         * Builds a new {@link Librispeech} object.
+         *
+         * @return the new {@link Librispeech} object
+         */
+        MRL getMrl() {
+            return repository.dataset(Application.Audio.ANY, groupId, artifactId, VERSION);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected Librispeech.Builder self() {
+            return this;
+        }
+    }
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/dataset/SpeechRecognitionDataset.java b/extensions/audio/src/main/java/ai/djl/audio/dataset/SpeechRecognitionDataset.java
new file mode 100644
index 000000000..29eefefff
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/dataset/SpeechRecognitionDataset.java
@@ -0,0 +1,176 @@
+/*
+ * 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.audio.dataset;
+
+import ai.djl.basicdataset.BasicDatasets;
+import ai.djl.basicdataset.utils.TextData;
+import ai.djl.engine.Engine;
+import ai.djl.modality.nlp.Vocabulary;
+import ai.djl.modality.nlp.embedding.EmbeddingException;
+import ai.djl.modality.nlp.embedding.TextEmbedding;
+import ai.djl.modality.nlp.embedding.TrainableWordEmbedding;
+import ai.djl.ndarray.NDManager;
+import ai.djl.repository.MRL;
+import ai.djl.repository.Repository;
+import ai.djl.training.dataset.RandomAccessDataset;
+import java.util.List;
+
+/**
+ * {@code SpeechRecognitionDataset} is an abstract dataset that can be used for datasets for
+ * Automatic Speech Recognition (ASR) where the source data is {@link AudioData} and the target data
+ * is {@link TextData}.
+ *
+ * <p>For the target data, it will create embeddings for the target data. Embeddings can be either
+ * pre-trained or trained on the go. Pre-trained {@link TextEmbedding} must be set in the {@code
+ * TextDataset.Builder}. If no embeddings are set, the dataset creates {@link
+ * TrainableWordEmbedding} based {@link TrainableWordEmbedding} from the {@link Vocabulary} created
+ * within the dataset.
+ *
+ * <p>For the source data, it will use the {@link ai.djl.audio.processor.AudioProcessor} to
+ * featurize data, if users want to write their own featurizer, they can get the original {@code
+ * NDArray} from {@link AudioData} without using any {@link ai.djl.audio.processor.AudioProcessor}.
+ */
+public abstract class SpeechRecognitionDataset extends RandomAccessDataset {
+
+    protected AudioData sourceAudioData;
+    protected TextData targetTextData;
+    protected NDManager manager;
+    protected Usage usage;
+
+    protected MRL mrl;
+    protected boolean prepared;
+
+    /**
+     * Creates a new instance of {@link SpeechRecognitionDataset} with the given necessary
+     * configurations.
+     *
+     * @param builder a builder with the necessary configurations
+     */
+    public SpeechRecognitionDataset(AudioBuilder<?> builder) {
+        super(builder);
+        sourceAudioData = new AudioData(AudioData.getDefaultConfiguration());
+        targetTextData = new TextData(TextData.getDefaultConfiguration());
+        manager = builder.manager;
+        usage = builder.usage;
+    }
+
+    /**
+     * Store and preprocess target text data.
+     *
+     * @param newTextData list of all unprocessed sentences in the dataset.
+     * @throws EmbeddingException if there is an error while embedding input.
+     */
+    protected void targetPreprocess(List<String> newTextData) throws EmbeddingException {
+        TextData textData = targetTextData;
+        textData.preprocess(
+                manager, newTextData.subList(0, (int) Math.min(limit, newTextData.size())));
+    }
+
+    /**
+     * This method is used to set the audio data path.
+     *
+     * @param audioPathList The path list of all original audio data
+     */
+    protected void sourcePreprocess(List<String> audioPathList) {
+        sourceAudioData.setAudioPaths(audioPathList);
+    }
+
+    /** Abstract AudioBuilder that helps build a {@link SpeechRecognitionDataset}. */
+    public abstract static class AudioBuilder<T extends AudioBuilder<T>> extends BaseBuilder<T> {
+
+        protected AudioData.Configuration sourceConfiguration;
+        protected TextData.Configuration targetConfiguration;
+        protected NDManager manager;
+
+        protected Repository repository;
+        protected String groupId;
+        protected Usage usage;
+        protected String artifactId;
+
+        /** Constructs a new builder. */
+        AudioBuilder() {
+            repository = BasicDatasets.REPOSITORY;
+            groupId = BasicDatasets.GROUP_ID;
+            usage = Usage.TRAIN;
+            sourceConfiguration = new AudioData.Configuration();
+            targetConfiguration = new TextData.Configuration();
+            manager = Engine.getInstance().newBaseManager();
+        }
+
+        /**
+         * Sets the {@link AudioData.Configuration} to use for the source text data.
+         *
+         * @param sourceConfiguration the {@link AudioData.Configuration}
+         * @return this builder
+         */
+        public T setSourceConfiguration(AudioData.Configuration sourceConfiguration) {
+            this.sourceConfiguration = sourceConfiguration;
+            return self();
+        }
+
+        /**
+         * Sets the {@link TextData.Configuration} to use for the target text data.
+         *
+         * @param targetConfiguration the {@link TextData.Configuration}
+         * @return this builder
+         */
+        public T setTargetConfiguration(TextData.Configuration targetConfiguration) {
+            this.targetConfiguration = targetConfiguration;
+            return self();
+        }
+
+        /**
+         * Sets the optional manager for the dataset (default follows engine default).
+         *
+         * @param manager the manager
+         * @return this builder
+         */
+        public T optManager(NDManager manager) {
+            this.manager = manager.newSubManager();
+            return self();
+        }
+
+        /**
+         * Sets the optional usage.
+         *
+         * @param usage the usage
+         * @return this builder
+         */
+        public T optUsage(Usage usage) {
+            this.usage = usage;
+            return self();
+        }
+
+        /**
+         * Sets the optional repository.
+         *
+         * @param repository the repository
+         * @return this builder
+         */
+        public T optRepository(Repository repository) {
+            this.repository = repository;
+            return self();
+        }
+
+        /**
+         * Sets optional groupId.
+         *
+         * @param groupId the groupId}
+         * @return this builder
+         */
+        public T optGroupId(String groupId) {
+            this.groupId = groupId;
+            return self();
+        }
+    }
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/dataset/package-info.java b/extensions/audio/src/main/java/ai/djl/audio/dataset/package-info.java
new file mode 100644
index 000000000..d2abdd35a
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/dataset/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 useful audio datasets. */
+package ai.djl.audio.dataset;
diff --git a/extensions/audio/src/main/java/ai/djl/audio/package-info.java b/extensions/audio/src/main/java/ai/djl/audio/package-info.java
new file mode 100644
index 000000000..d6322ae65
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/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 audio utils and datasets for djl. */
+package ai.djl.audio;
diff --git a/extensions/audio/src/main/java/ai/djl/audio/processor/AudioNormalizer.java b/extensions/audio/src/main/java/ai/djl/audio/processor/AudioNormalizer.java
new file mode 100644
index 000000000..55c4de4ed
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/processor/AudioNormalizer.java
@@ -0,0 +1,45 @@
+/*
+ * 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.audio.processor;
+
+import ai.djl.audio.AudioUtils;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+
+/** Use the mean and standard values to calculate the normalized values for audio signal. */
+public class AudioNormalizer implements AudioProcessor {
+
+    private static final float MAX_GAIN_DB = 300.0f;
+
+    private float targetDb;
+
+    /**
+     * Constructor for {@link AudioNormalizer}.
+     *
+     * @param targetDb target energy in decibels.
+     */
+    public AudioNormalizer(float targetDb) {
+        this.targetDb = targetDb;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray extractFeatures(NDManager manager, NDArray samples) {
+        float gain = targetDb - AudioUtils.rmsDb(samples);
+        gain = Math.min(gain, MAX_GAIN_DB);
+
+        float factor = (float) Math.pow(10f, gain / 20f);
+        samples = samples.mul(factor);
+        return samples;
+    }
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/processor/AudioProcessor.java b/extensions/audio/src/main/java/ai/djl/audio/processor/AudioProcessor.java
new file mode 100644
index 000000000..5c031b17b
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/processor/AudioProcessor.java
@@ -0,0 +1,29 @@
+/*
+ * 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.audio.processor;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+
+/** This interface is used for extracting features from origin audio samples. */
+public interface AudioProcessor {
+
+    /**
+     * Extracts features by the processor.
+     *
+     * @param manager The manager used for extracting features
+     * @param samples The Audio that needs to be extracting features
+     * @return Target feature
+     */
+    NDArray extractFeatures(NDManager manager, NDArray samples);
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/processor/LinearSpecgram.java b/extensions/audio/src/main/java/ai/djl/audio/processor/LinearSpecgram.java
new file mode 100644
index 000000000..cebf991b7
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/processor/LinearSpecgram.java
@@ -0,0 +1,150 @@
+/*
+ * 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.audio.processor;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDArrays;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import org.jtransforms.fft.FloatFFT_1D;
+
+/** Calculate linear spectrogram by short-time fourier transform. */
+public class LinearSpecgram implements AudioProcessor {
+
+    private static final float EPS = 1e-14f;
+
+    private float strideMs;
+    private float windowsMs;
+    private float sampleRate;
+
+    /**
+     * Calculate linear spectrogram by short-time fourier transform.
+     *
+     * @param strideMs Stride size of window
+     * @param windowsMs Window size
+     * @param sampleRate Sample rate of raw data
+     */
+    public LinearSpecgram(float strideMs, float windowsMs, int sampleRate) {
+        this.strideMs = strideMs;
+        this.windowsMs = windowsMs;
+        this.sampleRate = sampleRate;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray extractFeatures(NDManager manager, NDArray samples) {
+        return stft(samples);
+    }
+
+    private NDArray stft(NDArray samples) {
+        NDManager manager = samples.getManager();
+        int strideSize = (int) (0.001 * sampleRate * strideMs);
+        int windowSize = (int) (0.001 * sampleRate * windowsMs);
+        long truncateSize = (samples.size() - windowSize) % strideSize;
+        long len = samples.size() - truncateSize;
+        samples = samples.get(":" + len);
+
+        int rows = ((int) samples.size() - windowSize) / strideSize + 1;
+
+        NDList windowList = new NDList();
+        for (int row = 0; row < rows; row++) {
+            windowList.add(samples.get(strideSize * row + ":" + (strideSize * row + windowSize)));
+        }
+        samples = NDArrays.stack(windowList);
+
+        NDArray weighting = hanningWindow(windowSize, manager);
+        samples.muli(weighting);
+        NDList fftList = new NDList();
+        for (int row = 0; row < rows; row++) {
+            fftList.add(fft(samples.get(row)));
+        }
+
+        NDArray fft = NDArrays.stack(fftList).transpose();
+        fft = fft.pow(2);
+
+        weighting = weighting.pow(2);
+        NDArray scale = weighting.sum().mul(this.sampleRate);
+
+        NDArray middle = fft.get("1:-1,:");
+        middle = middle.mul(2).div(scale);
+        NDArray head = fft.get("0,:").div(scale).reshape(1, fft.getShape().get(1));
+        NDArray tail = fft.get("-1,:").div(scale).reshape(1, fft.getShape().get(1));
+        NDList list = new NDList(head, middle, tail);
+        fft = NDArrays.concat(list, 0);
+
+        NDArray freqsArray = manager.arange(fft.getShape().get(0));
+        freqsArray = freqsArray.mul(this.sampleRate / windowSize);
+
+        float[] freqs = freqsArray.toFloatArray();
+        int ind = 0;
+        for (int i = 0; i < freqs.length; i++) {
+            if (freqs[i] <= (this.sampleRate / 2)) {
+                ind = i;
+            } else {
+                break;
+            }
+        }
+        ind = ind + 1;
+
+        fft = fft.get(":" + ind + ",:").add(EPS);
+        fft = fft.log();
+
+        return fft;
+    }
+
+    private NDArray fft(NDArray in) {
+        float[] rawFFT = in.toFloatArray();
+        FloatFFT_1D fft = new FloatFFT_1D(rawFFT.length);
+        fft.realForward(rawFFT);
+        float[][] result;
+
+        int n = rawFFT.length;
+        if (n % 2 == 0) {
+            // n is even
+            result = new float[2][n / 2 + 1];
+            for (int i = 0; i < n / 2; i++) {
+                result[0][i] = rawFFT[2 * i]; // the real part fo the fast fourier transform
+                result[1][i] =
+                        rawFFT[2 * i + 1]; // the imaginary part of the fast fourier transform
+            }
+            result[1][0] = 0;
+            result[0][n / 2] = rawFFT[1];
+        } else {
+            // n is odd
+            result = new float[2][(n + 1) / 2];
+            for (int i = 0; i < n / 2; i++) {
+                result[0][i] = rawFFT[2 * i]; // the real part fo the fast fourier transform
+                result[1][i] =
+                        rawFFT[2 * i + 1]; // the imaginary part of the fast fourier transform
+            }
+            result[1][0] = 0;
+            result[1][(n - 1) / 2] = rawFFT[1];
+        }
+
+        float[] re = result[0]; // the real part fo the fast fourier transform
+        float[] im = result[1]; // the imaginary part of the fast fourier transform
+        float[] abs = new float[re.length];
+        for (int i = 0; i < re.length; i++) {
+            abs[i] = (float) Math.hypot(re[i], im[i]);
+        }
+        return in.getManager().create(abs);
+    }
+
+    private NDArray hanningWindow(int size, NDManager manager) {
+        float[] data = new float[size];
+        for (int i = 1; i < size; i++) {
+            data[i] = (float) (0.5 * (1 - Math.cos((2 * Math.PI * i) / (size - 1))));
+        }
+        return manager.create(data);
+    }
+}
diff --git a/extensions/audio/src/main/java/ai/djl/audio/processor/package-info.java b/extensions/audio/src/main/java/ai/djl/audio/processor/package-info.java
new file mode 100644
index 000000000..b3d688a12
--- /dev/null
+++ b/extensions/audio/src/main/java/ai/djl/audio/processor/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 multiple processors for audio and signal processing. */
+package ai.djl.audio.processor;
diff --git a/extensions/tablesaw/README.md b/extensions/tablesaw/README.md
new file mode 100644
index 000000000..f6901cc5a
--- /dev/null
+++ b/extensions/tablesaw/README.md
@@ -0,0 +1,28 @@
+# Tabular Dataset Support with Tablesaw
+
+This module contains the tabular dataset support with [Tablesaw](https://github.com/jtablesaw/tablesaw).
+
+The following functions have been implemented:
+
++ a `TablesawDataset` class extending `RandomAccessDataset` to support for importing tabular datasets in Tablesaw format.
+
+## Documentation
+
+You can build the latest javadocs locally using the following command:
+
+```sh
+./gradlew javadoc
+```
+The javadocs output is built in the `build/doc/javadoc` folder.
+
+## Installation
+
+You can pull the module from the central Maven repository by including the following dependency in your `pom.xml` file:
+
+```xml
+<dependency>
+    <groupId>ai.djl.tablesaw</groupId>
+    <artifactId>tablesaw</artifactId>
+    <version>0.18.0-SNAPSHOT</version>
+</dependency>
+```
diff --git a/extensions/tablesaw/build.gradle b/extensions/tablesaw/build.gradle
new file mode 100644
index 000000000..c3ed876da
--- /dev/null
+++ b/extensions/tablesaw/build.gradle
@@ -0,0 +1,30 @@
+group 'ai.djl.tablesaw'
+
+dependencies {
+    api project(":api")
+    api project(":basicdataset")
+    api "tech.tablesaw:tablesaw-core:${tablesaw_version}"
+    testImplementation("org.testng:testng:${testng_version}") {
+        exclude group: "junit", module: "junit"
+    }
+
+    testImplementation "org.slf4j:slf4j-simple:${slf4j_version}"
+    testImplementation project(":testing")
+
+//    testRuntimeOnly "tech.tablesaw:tablesaw-excel:${tablesaw_version}"
+//    testRuntimeOnly "tech.tablesaw:tablesaw-html:${tablesaw_version}"
+//    testRuntimeOnly "tech.tablesaw:tablesaw-json:${tablesaw_version}"
+    testRuntimeOnly project(":engines:pytorch:pytorch-engine")
+}
+
+publishing {
+    publications {
+        maven(MavenPublication) {
+            pom {
+                name = "Tablesaw toolkit for DJL"
+                description = "Tablesaw toolkit for DJL"
+                url = "http://www.djl.ai/extensions/${project.name}"
+            }
+        }
+    }
+}
diff --git a/extensions/tablesaw/gradlew b/extensions/tablesaw/gradlew
new file mode 100644
index 000000000..343e0d2ca
--- /dev/null
+++ b/extensions/tablesaw/gradlew
@@ -0,0 +1,1 @@
+../../gradlew
\ No newline at end of file
diff --git a/extensions/tablesaw/src/main/java/ai/djl/tablesaw/TablesawDataset.java b/extensions/tablesaw/src/main/java/ai/djl/tablesaw/TablesawDataset.java
new file mode 100644
index 000000000..874a4bfa3
--- /dev/null
+++ b/extensions/tablesaw/src/main/java/ai/djl/tablesaw/TablesawDataset.java
@@ -0,0 +1,239 @@
+/*
+ * 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.tablesaw;
+
+import ai.djl.basicdataset.utils.DynamicBuffer;
+import ai.djl.basicdataset.utils.Feature;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.training.dataset.RandomAccessDataset;
+import ai.djl.training.dataset.Record;
+import ai.djl.util.Progress;
+import java.nio.FloatBuffer;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import tech.tablesaw.api.Row;
+import tech.tablesaw.api.Table;
+import tech.tablesaw.io.ReadOptions;
+
+/** {@code TablesawDataset} represents the dataset that stored in a .csv file. */
+public class TablesawDataset extends RandomAccessDataset {
+
+    protected ReadOptions readOptions;
+    protected List<Feature> features;
+    protected List<Feature> labels;
+    protected Table table;
+
+    protected TablesawDataset(TablesawBuilder<?> builder) {
+        super(builder);
+        readOptions = builder.readOptions;
+        features = builder.features;
+        labels = builder.labels;
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public Record get(NDManager manager, long index) {
+        Row row = table.row(Math.toIntExact(index));
+        NDList data = toNDList(manager, row, features);
+        NDList label = toNDList(manager, row, labels);
+
+        return new Record(data, label);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    protected long availableSize() {
+        return table.rowCount();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void prepare(Progress progress) {
+        table = Table.read().usingOptions(readOptions);
+    }
+
+    /**
+     * Creates a builder to build a {@link TablesawDataset}.
+     *
+     * @return a new builder
+     */
+    public static TablesawBuilder<?> builder() {
+        return new TablesawBuilder<>();
+    }
+
+    /**
+     * Returns the column names of the Tablesaw file.
+     *
+     * @return a list of column name
+     */
+    public List<String> getColumnNames() {
+        if (table.isEmpty()) {
+            return Collections.emptyList();
+        }
+        return table.columnNames();
+    }
+
+    protected NDList toNDList(NDManager manager, Row row, List<Feature> selected) {
+        DynamicBuffer bb = new DynamicBuffer();
+        for (Feature feature : selected) {
+            String name = feature.getName();
+            String value = row.getString(name);
+            feature.getFeaturizer().featurize(bb, value);
+        }
+        FloatBuffer buf = bb.getBuffer();
+        return new NDList(manager.create(buf, new Shape(bb.getLength())));
+    }
+
+    /** Used to build a {@link TablesawDataset}. */
+    public static class TablesawBuilder<T extends TablesawBuilder<T>> extends BaseBuilder<T> {
+
+        protected ReadOptions readOptions;
+        protected List<Feature> features;
+        protected List<Feature> labels;
+
+        protected TablesawBuilder() {
+            features = new ArrayList<>();
+            labels = new ArrayList<>();
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        @SuppressWarnings("unchecked")
+        protected T self() {
+            return (T) this;
+        }
+
+        /**
+         * Sets the reading options.
+         *
+         * @param readOptions the {@code ReadOptions}
+         * @return this builder
+         */
+        public T setReadOptions(ReadOptions readOptions) {
+            this.readOptions = readOptions;
+            return self();
+        }
+
+        /**
+         * Adds the features to the feature set.
+         *
+         * @param features the features
+         * @return this builder
+         */
+        public T addFeature(Feature... features) {
+            Collections.addAll(this.features, features);
+            return self();
+        }
+
+        /**
+         * Adds a numeric feature to the feature set.
+         *
+         * @param name the feature name
+         * @return this builder
+         */
+        public T addNumericFeature(String name) {
+            features.add(new Feature(name, true));
+            return self();
+        }
+
+        /**
+         * Adds a categorical feature to the feature set.
+         *
+         * @param name the feature name
+         * @return this builder
+         */
+        public T addCategoricalFeature(String name) {
+            features.add(new Feature(name, false));
+            return self();
+        }
+
+        /**
+         * Adds a categorical feature to the feature set with specified mapping.
+         *
+         * @param name the feature name
+         * @param map a map contains categorical value maps to index
+         * @param onehotEncode true if use onehot encode
+         * @return this builder
+         */
+        public T addCategoricalFeature(
+                String name, Map<String, Integer> map, boolean onehotEncode) {
+            features.add(new Feature(name, map, onehotEncode));
+            return self();
+        }
+
+        /**
+         * Adds the features to the label set.
+         *
+         * @param labels the labels
+         * @return this builder
+         */
+        public T addLabel(Feature... labels) {
+            Collections.addAll(this.labels, labels);
+            return self();
+        }
+
+        /**
+         * Adds a number feature to the label set.
+         *
+         * @param name the label name
+         * @return this builder
+         */
+        public T addNumericLabel(String name) {
+            labels.add(new Feature(name, true));
+            return self();
+        }
+
+        /**
+         * Adds a categorical feature to the label set.
+         *
+         * @param name the feature name
+         * @return this builder
+         */
+        public T addCategoricalLabel(String name) {
+            labels.add(new Feature(name, true));
+            return self();
+        }
+
+        /**
+         * Adds a categorical feature to the feature set with specified mapping.
+         *
+         * @param name the feature name
+         * @param map a map contains categorical value maps to index
+         * @param onehotEncode true if use onehot encode
+         * @return this builder
+         */
+        public T addCategoricalLabel(String name, Map<String, Integer> map, boolean onehotEncode) {
+            labels.add(new Feature(name, map, onehotEncode));
+            return self();
+        }
+
+        /**
+         * Builds the new {@link TablesawDataset}.
+         *
+         * @return the new {@link TablesawDataset}
+         */
+        public TablesawDataset build() {
+            if (features.isEmpty()) {
+                throw new IllegalArgumentException("Missing features.");
+            }
+            if (labels.isEmpty()) {
+                throw new IllegalArgumentException("Missing labels.");
+            }
+            return new TablesawDataset(this);
+        }
+    }
+}
diff --git a/extensions/tablesaw/src/main/java/ai/djl/tablesaw/package-info.java b/extensions/tablesaw/src/main/java/ai/djl/tablesaw/package-info.java
new file mode 100644
index 000000000..2b2658a41
--- /dev/null
+++ b/extensions/tablesaw/src/main/java/ai/djl/tablesaw/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 classes to support tabular datasets in Tablesaw format. */
+package ai.djl.tablesaw;
diff --git a/extensions/tablesaw/src/main/javadoc/overview.html b/extensions/tablesaw/src/main/javadoc/overview.html
new file mode 100644
index 000000000..fbfd7c06f
--- /dev/null
+++ b/extensions/tablesaw/src/main/javadoc/overview.html
@@ -0,0 +1,10 @@
+<html>
+<head>
+    <meta charset="UTF-8">
+</head>
+<body>
+<p>This document is the API specification for the Tablesaw for DJL.</p>
+
+<p>The tablesaw module contains classes for DJL to support tabular datasets in Tablesaw format.</p>
+</body>
+</html>
diff --git a/gradle.properties b/gradle.properties
index 0ff6242d6..64cc4d26c 100644
--- a/gradle.properties
+++ b/gradle.properties
@@ -19,7 +19,8 @@ paddlepaddle_version=2.2.2
 sentencepiece_version=0.1.95
 tokenizers_version=0.11.0
 fasttext_version=0.9.2
-xgboost_version=1.6.0
+xgboost_version=1.6.1
+rapis_version=22.04.0
 
 commons_cli_version=1.5.0
 commons_compress_version=1.21
@@ -33,6 +34,7 @@ awssdk_version=2.17.151
 hadoop_version=3.3.1
 javacpp_version=1.5.6
 protobuf_version=3.19.2
+tablesaw_version=0.43.1
 
 antlr_version=4.9.3
 testng_version=7.5
diff --git a/settings.gradle b/settings.gradle
index 25492c51b..859c5362e 100644
--- a/settings.gradle
+++ b/settings.gradle
@@ -25,6 +25,7 @@ include ':engines:tensorrt'
 include ':engines:tflite:tflite-engine'
 include ':engines:tflite:tflite-native'
 include ':examples'
+include 'extensions:audio'
 include ':extensions:aws-ai'
 include ':extensions:benchmark'
 include ':extensions:fasttext'
@@ -32,7 +33,9 @@ include ':extensions:hadoop'
 include ':extensions:opencv'
 include ':extensions:sentencepiece'
 include ':extensions:tokenizers'
+include ':extensions:tablesaw'
 include ':integration'
 include ':model-zoo'
 include ':testing'
 
+
diff --git a/tools/gradle/publish.gradle b/tools/gradle/publish.gradle
index aff4d8de2..1c02921e1 100644
--- a/tools/gradle/publish.gradle
+++ b/tools/gradle/publish.gradle
@@ -16,11 +16,13 @@ configure([
         project(':engines:tensorflow:tensorflow-model-zoo'),
         project(':engines:tensorrt'),
         project(':engines:tflite:tflite-engine'),
+        project(':extensions:audio'),
         project(':extensions:aws-ai'),
         project(':extensions:fasttext'),
         project(':extensions:hadoop'),
         project(':extensions:opencv'),
         project(':extensions:sentencepiece'),
+        project(':extensions:tablesaw'),
         project(':extensions:tokenizers'),
         project(':model-zoo')]) {
     apply plugin: "maven-publish"


Test patch:
diff --git a/api/src/test/java/ai/djl/nn/norm/GhostBatchNormTest.java b/api/src/test/java/ai/djl/nn/norm/GhostBatchNormTest.java
new file mode 100644
index 000000000..3f33a1242
--- /dev/null
+++ b/api/src/test/java/ai/djl/nn/norm/GhostBatchNormTest.java
@@ -0,0 +1,52 @@
+/*
+ * 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.nn.norm;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class GhostBatchNormTest {
+
+    @Test
+    public void originalBatchOfSizeOneSubBatchIntoOneVbsListOfLengthOne() {
+        testSubBatchingShapeSize(new Shape(1, 1, 10), 1, 1);
+    }
+
+    @Test
+    public void originalBatchOfSizeSixSubBatchIntoTwoVbsListOfLengthThree() {
+        testSubBatchingShapeSize(new Shape(6, 1, 10), 2, 3);
+    }
+
+    @Test
+    public void originalBatchOfSize60SubBatchInto64VbsListOfLengthOne() {
+        testSubBatchingShapeSize(new Shape(60, 10), 64, 1);
+    }
+
+    private void testSubBatchingShapeSize(Shape shape, int vbs, int expectedSize) {
+        NDList[] zeroSubList = generateZerosArrayAndSubBatch(shape, vbs);
+        Assert.assertEquals(expectedSize, zeroSubList.length);
+    }
+
+    private NDList[] generateZerosArrayAndSubBatch(Shape shape, int vbs) {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            GhostBatchNorm gbn =
+                    new GhostBatchNorm(GhostBatchNorm.builder().optVirtualBatchSize(vbs));
+            NDArray zerosInput = manager.zeros(shape);
+            return gbn.split(new NDList(zerosInput));
+        }
+    }
+}
diff --git a/api/src/test/java/ai/djl/nn/norm/package-info.java b/api/src/test/java/ai/djl/nn/norm/package-info.java
new file mode 100644
index 000000000..aed4ebedf
--- /dev/null
+++ b/api/src/test/java/ai/djl/nn/norm/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 tests for {@link ai.djl.nn.norm}. */
+package ai.djl.nn.norm;
diff --git a/basicdataset/src/test/java/ai/djl/basicdataset/UniversalDependenciesEnglishEWTTest.java b/basicdataset/src/test/java/ai/djl/basicdataset/UniversalDependenciesEnglishEWTTest.java
new file mode 100644
index 000000000..f6764be88
--- /dev/null
+++ b/basicdataset/src/test/java/ai/djl/basicdataset/UniversalDependenciesEnglishEWTTest.java
@@ -0,0 +1,110 @@
+/*
+ * 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.basicdataset;
+
+import ai.djl.basicdataset.nlp.UniversalDependenciesEnglishEWT;
+import ai.djl.basicdataset.utils.TextData;
+import ai.djl.basicdataset.utils.TextData.Configuration;
+import ai.djl.modality.nlp.preprocess.LowerCaseConvertor;
+import ai.djl.modality.nlp.preprocess.SimpleTokenizer;
+import ai.djl.ndarray.NDManager;
+import ai.djl.training.dataset.Dataset;
+import ai.djl.training.dataset.Record;
+import ai.djl.translate.TranslateException;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Locale;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class UniversalDependenciesEnglishEWTTest {
+
+    private static final int EMBEDDING_SIZE = 15;
+
+    @Test
+    public void testGetDataWithPreTrainedEmbedding() throws IOException, TranslateException {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            UniversalDependenciesEnglishEWT universalDependenciesEnglishEWT =
+                    UniversalDependenciesEnglishEWT.builder()
+                            .optUsage(Dataset.Usage.TRAIN)
+                            .setSourceConfiguration(
+                                    new Configuration()
+                                            .setTextProcessors(
+                                                    Arrays.asList(
+                                                            new SimpleTokenizer(),
+                                                            new LowerCaseConvertor(Locale.ENGLISH)))
+                                            .setTextEmbedding(
+                                                    TestUtils.getTextEmbedding(
+                                                            manager, EMBEDDING_SIZE)))
+                            .setSampling(32, true)
+                            .optLimit(10)
+                            .build();
+
+            universalDependenciesEnglishEWT.prepare();
+            Record record = universalDependenciesEnglishEWT.get(manager, 0);
+            Assert.assertEquals(record.getData().get(0).getShape().dimension(), 2);
+            Assert.assertEquals(record.getLabels().get(0).getShape().dimension(), 1);
+        }
+    }
+
+    @Test
+    public void testGetDataWithTrainableEmbedding() throws IOException, TranslateException {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            UniversalDependenciesEnglishEWT universalDependenciesEnglishEWT =
+                    UniversalDependenciesEnglishEWT.builder()
+                            .optUsage(Dataset.Usage.TEST)
+                            .setSourceConfiguration(
+                                    new Configuration()
+                                            .setTextProcessors(
+                                                    Arrays.asList(
+                                                            new SimpleTokenizer(),
+                                                            new LowerCaseConvertor(Locale.ENGLISH)))
+                                            .setEmbeddingSize(EMBEDDING_SIZE))
+                            .setSampling(32, true)
+                            .optLimit(10)
+                            .build();
+
+            universalDependenciesEnglishEWT.prepare();
+            Record record = universalDependenciesEnglishEWT.get(manager, 0);
+            Assert.assertEquals(record.getData().size(), 1);
+            Assert.assertEquals(record.getData().get(0).getShape().dimension(), 1);
+            Assert.assertEquals(record.getLabels().size(), 1);
+            Assert.assertEquals(record.getLabels().get(0).getShape().dimension(), 1);
+        }
+    }
+
+    @Test
+    public void testMisc() throws TranslateException, IOException {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            UniversalDependenciesEnglishEWT universalDependenciesEnglishEWT =
+                    UniversalDependenciesEnglishEWT.builder()
+                            .optUsage(Dataset.Usage.VALIDATION)
+                            .setSourceConfiguration(
+                                    new TextData.Configuration()
+                                            .setTextProcessors(
+                                                    Arrays.asList(
+                                                            new SimpleTokenizer(),
+                                                            new LowerCaseConvertor(Locale.ENGLISH)))
+                                            .setTextEmbedding(
+                                                    TestUtils.getTextEmbedding(
+                                                            manager, EMBEDDING_SIZE)))
+                            .setSampling(32, true)
+                            .optLimit(350)
+                            .build();
+
+            universalDependenciesEnglishEWT.prepare();
+            universalDependenciesEnglishEWT.prepare();
+            Assert.assertEquals(universalDependenciesEnglishEWT.size(), 350);
+        }
+    }
+}
diff --git a/basicdataset/src/test/resources/mlrepo/dataset/nlp/ai/djl/basicdataset/universal-dependencies/universal-dependencies-en-ewt/metadata.json b/basicdataset/src/test/resources/mlrepo/dataset/nlp/ai/djl/basicdataset/universal-dependencies/universal-dependencies-en-ewt/metadata.json
new file mode 100644
index 000000000..46719de26
--- /dev/null
+++ b/basicdataset/src/test/resources/mlrepo/dataset/nlp/ai/djl/basicdataset/universal-dependencies/universal-dependencies-en-ewt/metadata.json
@@ -0,0 +1,29 @@
+{
+  "metadataVersion": "0.2",
+  "resourceType": "dataset",
+  "application": "nlp",
+  "groupId": "ai.djl.basicdataset.universal-dependencies",
+  "artifactId": "universal-dependencies-en-ewt",
+  "name": "universal-dependencies-en-ewt",
+  "description": "A Gold Standard Universal Dependencies Corpus for English, built over the source material of the English Web Treebank LDC2012T13.",
+  "website": "https://universaldependencies.org/treebanks/en_ewt/index.html",
+  "licenses": {
+    "license": {
+      "name": "Creative Commons Attribution-ShareAlike 4.0 International (CC BY-SA 4.0)",
+      "url": "https://creativecommons.org/licenses/by-sa/4.0/"
+    }
+  },
+  "artifacts": [
+    {
+      "version": "2.0",
+      "snapshot": false,
+      "files": {
+        "universal-dependencies-en-ewt": {
+          "uri": "https://bitbucket.org/sivareddyg/public/downloads/en-ud-v2.zip",
+          "sha1Hash": "27a39d8d2686c9945842d26c53c2e51d1ce48226",
+          "size": 688379
+        }
+      }
+    }
+  ]
+}
diff --git a/engines/ml/xgboost/src/test/java/ai/djl/ml/xgboost/XgbModelTest.java b/engines/ml/xgboost/src/test/java/ai/djl/ml/xgboost/XgbModelTest.java
index 4310b3426..c29853e90 100644
--- a/engines/ml/xgboost/src/test/java/ai/djl/ml/xgboost/XgbModelTest.java
+++ b/engines/ml/xgboost/src/test/java/ai/djl/ml/xgboost/XgbModelTest.java
@@ -15,6 +15,7 @@ package ai.djl.ml.xgboost;
 
 import ai.djl.MalformedModelException;
 import ai.djl.Model;
+import ai.djl.engine.Engine;
 import ai.djl.inference.Predictor;
 import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
@@ -45,6 +46,30 @@ public class XgbModelTest {
                 modelDir.resolve("regression.json").toString());
     }
 
+    @Test
+    public void testVersion() {
+        Engine engine = Engine.getEngine("XGBoost");
+        Assert.assertEquals("1.6.1", engine.getVersion());
+    }
+
+    /*
+    This test depends on rapis, which doesn't work for CPU
+    @Test
+    public void testCudf() {
+        Engine engine = Engine.getEngine("XGBoost");
+        if (engine.hasCapability(StandardCapabilities.CUDA)) {
+            Float[] data = new Float[]{1.f, null, 5.f, 7.f, 9.f};
+            Table x = new Table.TestBuilder().column(data).build();
+            ColumnBatch columnBatch = new CudfColumnBatch(x, null, null, null);
+
+            try (XgbNDManager manager = (XgbNDManager)engine.newBaseManager()) {
+                NDArray array = manager.create(columnBatch,0, 1);
+                Assert.assertEquals(array.getShape().get(0), 5);
+            }
+        }
+    }
+    */
+
     @Test
     public void testLoad() throws MalformedModelException, IOException, TranslateException {
         // TODO: skip for Windows since it is not supported by XGBoost
diff --git a/extensions/audio/src/test/java/ai/djl/audio/dataset/LibrispeechTest.java b/extensions/audio/src/test/java/ai/djl/audio/dataset/LibrispeechTest.java
new file mode 100644
index 000000000..5f776a1d6
--- /dev/null
+++ b/extensions/audio/src/test/java/ai/djl/audio/dataset/LibrispeechTest.java
@@ -0,0 +1,47 @@
+/*
+ * 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.audio.dataset;
+
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.training.dataset.Dataset;
+import ai.djl.translate.TranslateException;
+import java.io.IOException;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class LibrispeechTest {
+
+    @Test
+    public void testLibrispeech() throws IOException, TranslateException {
+        // The dataset is very large, only run the test when -Dweekly=true is set
+        if (Boolean.getBoolean("weekly")) {
+            NDManager manager = NDManager.newBaseManager();
+            Librispeech dataset =
+                    Librispeech.builder()
+                            .optUsage(Dataset.Usage.TEST)
+                            .setSampling(32, false)
+                            .build();
+            dataset.prepare();
+            NDList data = dataset.get(manager, 0).getData();
+            Assert.assertEquals(data.getShapes()[0].getShape(), new long[] {161, 306});
+
+            long[] expected = {
+                0, 1, 2, 3, 2344, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 2, 14, 15, 16, 5, 2344, 17, 18,
+                19
+            };
+            NDList labels = dataset.get(manager, 0).getLabels();
+            Assert.assertEquals(labels.get(0).toLongArray(), expected);
+        }
+    }
+}
diff --git a/extensions/audio/src/test/java/ai/djl/audio/dataset/package-info.java b/extensions/audio/src/test/java/ai/djl/audio/dataset/package-info.java
new file mode 100644
index 000000000..74110d1c7
--- /dev/null
+++ b/extensions/audio/src/test/java/ai/djl/audio/dataset/package-info.java
@@ -0,0 +1,16 @@
+/*
+ * 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 the dataset module. */
+package ai.djl.audio.dataset;
\ No newline at end of file
diff --git a/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java b/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java
new file mode 100644
index 000000000..28cec1b62
--- /dev/null
+++ b/extensions/audio/src/test/java/ai/djl/audio/processor/AudioProcessorTest.java
@@ -0,0 +1,67 @@
+/*
+ * 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.audio.processor;
+
+import ai.djl.Device;
+import ai.djl.audio.AudioUtils;
+import ai.djl.audio.dataset.AudioData;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.training.util.DownloadUtils;
+import java.io.IOException;
+import java.util.Arrays;
+import java.util.Collections;
+import org.testng.Assert;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+public class AudioProcessorTest {
+
+    private static final String URL = "https://resources.djl.ai/audios/test_01.wav";
+
+    @BeforeClass
+    public void setUp() throws IOException {
+        DownloadUtils.download(URL, "build/test/test_01.wav");
+    }
+
+    @Test
+    public void testAudioNormalizer() {
+        NDManager manager = NDManager.newBaseManager(Device.cpu());
+        AudioData.Configuration configuration =
+                new AudioData.Configuration()
+                        .setProcessorList(Collections.singletonList(new AudioNormalizer(-20)));
+        AudioData testData = new AudioData(configuration);
+        testData.setAudioPaths(Collections.singletonList("build/test/test_01.wav"));
+        NDArray samples = testData.getPreprocessedData(manager, 0);
+        Assert.assertEquals(AudioUtils.rmsDb(samples), -20.0f, 1e-3);
+    }
+
+    @Test
+    public static void testLinearSpecgram() {
+        NDManager manager = NDManager.newBaseManager(Device.cpu());
+        AudioData.Configuration configuration =
+                new AudioData.Configuration()
+                        .setSampleRate(16000)
+                        .setProcessorList(
+                                Arrays.asList(
+                                        new AudioNormalizer(-20),
+                                        new LinearSpecgram(10, 20, 16000)));
+        AudioData testData = new AudioData(configuration);
+        testData.setAudioPaths(Collections.singletonList("build/test/test_01.wav"));
+        NDArray samples = testData.getPreprocessedData(manager, 0);
+        Assert.assertEquals(new Shape(161, 838), samples.getShape());
+        Assert.assertEquals(samples.get("0,0").toFloatArray()[0], -15.4571f, 1e-3f);
+    }
+}
diff --git a/extensions/audio/src/test/java/ai/djl/audio/processor/package-info.java b/extensions/audio/src/test/java/ai/djl/audio/processor/package-info.java
new file mode 100644
index 000000000..3b8fc9e92
--- /dev/null
+++ b/extensions/audio/src/test/java/ai/djl/audio/processor/package-info.java
@@ -0,0 +1,16 @@
+/*
+ * 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 the processor module. */
+package ai.djl.audio.processor;
\ No newline at end of file
diff --git a/extensions/audio/src/test/resources/mlrepo/dataset/audio/ai/djl/basicdataset/librispeech/metadata.json b/extensions/audio/src/test/resources/mlrepo/dataset/audio/ai/djl/basicdataset/librispeech/metadata.json
new file mode 100644
index 000000000..5224958f9
--- /dev/null
+++ b/extensions/audio/src/test/resources/mlrepo/dataset/audio/ai/djl/basicdataset/librispeech/metadata.json
@@ -0,0 +1,35 @@
+{
+  "metadataVersion": "0.2",
+  "resourceType": "dataset",
+  "application": "audio",
+  "groupId": "ai.djl.basicdata",
+  "artifactId": "librispeech",
+  "name": "librispeech",
+  "description": "LibriSpeech is a corpus of approximately 1000 hours of 16kHz read English speech, prepared by Vassil Panayotov with the assistance of Daniel Povey. The data is derived from read audiobooks from the LibriVox project, and has been carefully segmented and aligned.",
+  "website": "https://huggingface.co/datasets/librispeech_asr",
+  "licenses": {
+    "license": {
+      "name": "CC BY 4.0",
+      "url": "https://creativecommons.org/licenses/by/4.0/"
+    }
+  },
+  "artifacts": [
+    {
+      "version": "1.0",
+      "snapshot": false,
+      "name": "Librispeech",
+      "files": {
+        "train": {
+          "uri": "http://us.openslr.org/resources/12/train-clean-100.tar.gz",
+          "sha1Hash": "cf83135102b2958182e26ef3b554b439338aa48e",
+          "size": 6387309499
+        },
+        "test": {
+          "uri": "http://us.openslr.org/resources/12/test-clean.tar.gz",
+          "sha1Hash": "2978149c5fa469821a02171e1d547cf1065f1213",
+          "size": 346663984
+        }
+      }
+    }
+  ]
+}
diff --git a/extensions/tablesaw/src/test/java/ai/djl/tablesaw/TablesawDatasetTest.java b/extensions/tablesaw/src/test/java/ai/djl/tablesaw/TablesawDatasetTest.java
new file mode 100644
index 000000000..be7d71dd4
--- /dev/null
+++ b/extensions/tablesaw/src/test/java/ai/djl/tablesaw/TablesawDatasetTest.java
@@ -0,0 +1,72 @@
+/*
+ * 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.tablesaw;
+
+import ai.djl.Model;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.nn.Blocks;
+import ai.djl.training.dataset.Record;
+import ai.djl.translate.TranslateException;
+import java.io.IOException;
+import java.net.URL;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+import tech.tablesaw.api.ColumnType;
+import tech.tablesaw.io.csv.CsvReadOptions;
+
+public class TablesawDatasetTest {
+
+    @Test
+    public void testTablesawDataset() throws IOException, TranslateException {
+        URL dailyDelhiClimateTest =
+                new URL(
+                        "https://mlrepo.djl.ai/dataset/tabular/ai/djl/basicdataset/daily-delhi-climate/3.0/DailyDelhiClimateTest.csv");
+
+        try (Model model = Model.newInstance("model")) {
+            model.setBlock(Blocks.identityBlock());
+
+            NDManager manager = model.getNDManager();
+            TablesawDataset dailyDelhiClimate =
+                    TablesawDataset.builder()
+                            .setReadOptions(
+                                    CsvReadOptions.builder(dailyDelhiClimateTest)
+                                            .header(true)
+                                            .columnTypes(
+                                                    new ColumnType[] {
+                                                        ColumnType.STRING,
+                                                        ColumnType.STRING,
+                                                        ColumnType.STRING,
+                                                        ColumnType.STRING,
+                                                        ColumnType.STRING
+                                                    })
+                                            .build())
+                            .addNumericFeature("meantemp")
+                            .addNumericFeature("meanpressure")
+                            .addNumericLabel("humidity")
+                            .setSampling(32, true)
+                            .build();
+            dailyDelhiClimate.prepare();
+
+            long size = dailyDelhiClimate.size();
+            Assert.assertEquals(size, 114);
+
+            Record record = dailyDelhiClimate.get(manager, 3);
+            NDList data = record.getData();
+
+            NDList labels = record.getLabels();
+            Assert.assertEquals(data.head().toFloatArray(), new float[] {18.7f, 1015.7f});
+            Assert.assertEquals(labels.head().toFloatArray(), new float[] {70.05f});
+        }
+    }
+}
diff --git a/extensions/tablesaw/src/test/java/ai/djl/tablesaw/package-info.java b/extensions/tablesaw/src/test/java/ai/djl/tablesaw/package-info.java
new file mode 100644
index 000000000..8f1d80741
--- /dev/null
+++ b/extensions/tablesaw/src/test/java/ai/djl/tablesaw/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 tests for the tablesaw module. */
+package ai.djl.tablesaw;
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayAttachmentTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayAttachmentTest.java
new file mode 100644
index 000000000..e83bceb05
--- /dev/null
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDArrayAttachmentTest.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.integration.tests.ndarray;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.index.NDIndex;
+import ai.djl.ndarray.types.Shape;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class NDArrayAttachmentTest {
+
+    @Test
+    public void testReturnResource() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray array3x4 = manager.ones(new Shape(3, 4));
+            try (NDManager subManager = NDManager.newBaseManager()) {
+                array3x4.tempAttach(subManager);
+                Assert.assertEquals(array3x4.getManager(), subManager);
+            }
+            Assert.assertEquals(array3x4.getManager(), manager);
+        }
+    }
+
+    @Test
+    public void testIndexationUsesSpecificManager() {
+        try (NDManager manager = NDManager.newBaseManager()) {
+            NDArray array3x4 = manager.ones(new Shape(3, 4));
+            array3x4.setName("Test()");
+            NDArray array4 = array3x4.get(1);
+            Assert.assertEquals(array4.getManager(), manager);
+            try (NDManager subManager = NDManager.newBaseManager()) {
+                NDArray array4sub1 = array3x4.get(subManager, 1);
+                Assert.assertEquals(array4sub1.getManager(), subManager);
+                NDArray array4sub2 = array3x4.get(subManager, new NDIndex(1));
+                Assert.assertEquals(array4sub2.getManager(), subManager);
+            }
+        }
+    }
+}
diff --git a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
index 49e9f652c..88abee3ad 100644
--- a/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/ndarray/NDIndexTest.java
@@ -55,9 +55,9 @@ public class NDIndexTest {
     @Test
     public void testGather() {
         try (NDManager manager = NDManager.newBaseManager()) {
-            NDArray arr = manager.arange(20f).reshape(-1, 4);
+            NDArray original = manager.arange(20f).reshape(-1, 4);
             NDArray index = manager.create(new long[] {0, 0, 2, 1, 1, 2}, new Shape(3, 2));
-            NDArray actual = arr.gather(index, 1);
+            NDArray actual = original.gather(index, 1);
             NDArray expected = manager.create(new float[] {0, 0, 6, 5, 9, 10}, new Shape(3, 2));
             Assert.assertEquals(actual, expected);
         }
@@ -66,9 +66,9 @@ public class NDIndexTest {
     @Test
     public void testTake() {
         try (NDManager manager = NDManager.newBaseManager()) {
-            NDArray arr = manager.arange(1, 7f).reshape(-1, 3);
+            NDArray original = manager.arange(1, 7f).reshape(-1, 3);
             NDArray index = manager.create(new long[] {0, 4, 1, 2}, new Shape(2, 2));
-            NDArray actual = arr.take(index);
+            NDArray actual = original.take(index);
             NDArray expected = manager.create(new float[] {1, 5, 2, 3}, new Shape(2, 2));
             Assert.assertEquals(actual, expected);
         }
diff --git a/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java b/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
index f8f53ddbf..a30f8f382 100644
--- a/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/nn/BlockCoreTest.java
@@ -36,6 +36,7 @@ import ai.djl.nn.convolutional.Conv3d;
 import ai.djl.nn.core.Linear;
 import ai.djl.nn.norm.BatchNorm;
 import ai.djl.nn.norm.Dropout;
+import ai.djl.nn.norm.GhostBatchNorm;
 import ai.djl.nn.norm.LayerNorm;
 import ai.djl.nn.recurrent.GRU;
 import ai.djl.nn.recurrent.LSTM;
@@ -204,6 +205,41 @@ public class BlockCoreTest {
         }
     }
 
+    @SuppressWarnings("try")
+    @Test
+    public void testGhostBatchNorm() throws IOException, MalformedModelException {
+        TrainingConfig config =
+                new DefaultTrainingConfig(Loss.l2Loss())
+                        .optInitializer(Initializer.ONES, Parameter.Type.WEIGHT);
+
+        Block block = GhostBatchNorm.builder().build();
+        try (Model model = Model.newInstance("model")) {
+            model.setBlock(block);
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                // the unused GradientCollector is for BatchNorm to know it is on training mode
+                try (GradientCollector collector = trainer.newGradientCollector()) {
+                    Shape inputShape = new Shape(1, 1, 10);
+                    trainer.initialize(inputShape);
+
+                    NDManager manager = trainer.getManager();
+                    NDArray data = manager.arange(10f).reshape(inputShape);
+                    NDArray expected =
+                            manager.create(
+                                    new float[] {
+                                        -1.5667f, -1.2185f, -0.8704f, -0.5222f, -0.1741f, 0.1741f,
+                                        0.5222f, 0.8704f, 1.2185f, 1.5667f
+                                    },
+                                    inputShape);
+                    NDArray result = trainer.forward(new NDList(data)).singletonOrThrow();
+
+                    Assertions.assertAlmostEquals(result, expected);
+                    testEncode(manager, block);
+                }
+            }
+        }
+    }
+
     @SuppressWarnings("try")
     @Test
     public void testLayerNorm() throws IOException, MalformedModelException {