Repository: deepjavalibrary/djl
Pull request #2057 (2022-10-02 07:56:35)

TabNet Training
1. checkout the bugs in TabNet and fix them
2. change tabNetLoss as an embedding mode(add L2Loss for it)
3. add Tabular for tabular training in djl-Zero
4. add an example on airfoil dataset on how tabNet should be used and its corresponding test


codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2057?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) Report
Base: **72.08**% // Head: **69.70**% // Decreases project coverage by **`-2.38%`** :warning:
> Coverage data is based on head [(`c6f889d`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2057?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary) compared to base [(`bb5073f`)](https://codecov.io/gh/deepjavalibrary/djl/commit/bb5073f7ffc33f1883621fc1783dbb01e70b418f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary).
> Patch coverage: 68.91% of modified lines in pull request are covered.

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


```diff
@@             Coverage Diff              @@
##             master    #2057      +/-   ##
============================================
- Coverage     72.08%   69.70%   -2.39%     
- Complexity     5126     5981     +855     
============================================
  Files           473      600     +127     
  Lines         21970    26540    +4570     
  Branches       2351     2883     +532     
============================================
+ Hits          15838    18499    +2661     
- Misses         4925     6656    +1731     
- Partials       1207     1385     +178     
```


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

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

</details>

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


Patch:
diff --git a/api/src/main/java/ai/djl/training/loss/TabNetLoss.java b/api/src/main/java/ai/djl/training/loss/TabNetClassificationLoss.java
similarity index 53%
rename from api/src/main/java/ai/djl/training/loss/TabNetLoss.java
rename to api/src/main/java/ai/djl/training/loss/TabNetClassificationLoss.java
index 296f4b93d..58756054d 100644
--- a/api/src/main/java/ai/djl/training/loss/TabNetLoss.java
+++ b/api/src/main/java/ai/djl/training/loss/TabNetClassificationLoss.java
@@ -16,29 +16,32 @@ import ai.djl.ndarray.NDArray;
 import ai.djl.ndarray.NDList;
 
 /**
- * Calculates the loss for tabNet. Actually, loss has been calculated through the forward function
- * of tabNet. What's done here is just getting the loss function from prediction.
+ * Calculates the loss for tabNet in Classification tasks.
+ *
+ * <p>Actually, tabNet is not only used for Supervised Learning, it's also widely used in
+ * unsupervised learning. For unsupervised learning, it should come from the decoder(aka
+ * attentionTransformer of tabNet)
  */
-public class TabNetLoss extends Loss {
-    /** Calculates the loss of a TabNet instance. */
-    public TabNetLoss() {
-        this("TabNetLoss");
+public final class TabNetClassificationLoss extends Loss {
+    /** Calculates the loss of a TabNet instance for regression tasks. */
+    public TabNetClassificationLoss() {
+        this("TabNetClassificationLoss");
     }
 
     /**
-     * Calculates the loss of a TabNet instance.
+     * Calculates the loss of a TabNet instance for regression tasks.
      *
      * @param name the name of the loss function
      */
-    public TabNetLoss(String name) {
+    public TabNetClassificationLoss(String name) {
         super(name);
     }
 
     /** {@inheritDoc} */
     @Override
     public NDArray evaluate(NDList labels, NDList predictions) {
-        // loss is already calculated inside the forward of tabNet
-        // so here we just need to get it out from prediction
-        return predictions.get(1);
+        return Loss.softmaxCrossEntropyLoss()
+                .evaluate(labels, new NDList(predictions.get(0)))
+                .add(predictions.get(1));
     }
 }
diff --git a/api/src/main/java/ai/djl/training/loss/TabNetRegressionLoss.java b/api/src/main/java/ai/djl/training/loss/TabNetRegressionLoss.java
new file mode 100644
index 000000000..21f290136
--- /dev/null
+++ b/api/src/main/java/ai/djl/training/loss/TabNetRegressionLoss.java
@@ -0,0 +1,51 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.training.loss;
+
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+
+/**
+ * Calculates the loss of tabNet for regression tasks.
+ *
+ * <p>Actually, tabNet is not only used for Supervised Learning, it's also widely used in
+ * unsupervised learning. For unsupervised learning, it should come from the decoder(aka
+ * attentionTransformer of tabNet)
+ */
+public class TabNetRegressionLoss extends Loss {
+    /** Calculates the loss of a TabNet instance for regression tasks. */
+    public TabNetRegressionLoss() {
+        this("TabNetRegressionLoss");
+    }
+
+    /**
+     * Calculates the loss of a TabNet instance for regression tasks.
+     *
+     * @param name the name of the loss function
+     */
+    public TabNetRegressionLoss(String name) {
+        super(name);
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public NDArray evaluate(NDList labels, NDList predictions) {
+        // sparseLoss is already calculated inside the forward of tabNet
+        // so here we just need to get it out from prediction
+        return labels.singletonOrThrow()
+                .sub(predictions.get(0))
+                .square()
+                .mean()
+                .add(predictions.get(1));
+    }
+}
diff --git a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/TabularDataset.java b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/TabularDataset.java
index 43ff98dc7..9a8380c1c 100644
--- a/basicdataset/src/main/java/ai/djl/basicdataset/tabular/TabularDataset.java
+++ b/basicdataset/src/main/java/ai/djl/basicdataset/tabular/TabularDataset.java
@@ -46,6 +46,24 @@ public abstract class TabularDataset extends RandomAccessDataset {
         labels = builder.labels;
     }
 
+    /**
+     * Gets the feature size of current {@link TabularDataset}.
+     *
+     * @return the feature size
+     */
+    public int getFeatureSize() {
+        return features.size();
+    }
+
+    /**
+     * Gets the label size of current {@link TabularDataset}.
+     *
+     * @return the feature size
+     */
+    public int getLabelSize() {
+        return labels.size();
+    }
+
     /** {@inheritDoc} */
     @Override
     public Record get(NDManager manager, long index) {
diff --git a/djl-zero/src/main/java/ai/djl/zero/tabular/Tabular.java b/djl-zero/src/main/java/ai/djl/zero/tabular/Tabular.java
new file mode 100644
index 000000000..634466a0d
--- /dev/null
+++ b/djl-zero/src/main/java/ai/djl/zero/tabular/Tabular.java
@@ -0,0 +1,97 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.zero.tabular;
+
+import ai.djl.Model;
+import ai.djl.basicdataset.tabular.TabularDataset;
+import ai.djl.basicmodelzoo.tabular.TabNet;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Block;
+import ai.djl.repository.zoo.ZooModel;
+import ai.djl.training.DefaultTrainingConfig;
+import ai.djl.training.EasyTrain;
+import ai.djl.training.Trainer;
+import ai.djl.training.TrainingConfig;
+import ai.djl.training.dataset.Dataset;
+import ai.djl.training.listener.TrainingListener;
+import ai.djl.training.loss.TabNetRegressionLoss;
+import ai.djl.translate.NoopTranslator;
+import ai.djl.translate.TranslateException;
+import ai.djl.zero.Performance;
+
+import java.io.IOException;
+
+/** Tabular takes a NDList as input and output an NDList (for supervised learning). */
+public final class Tabular {
+
+    private Tabular() {}
+
+    /**
+     * Trains a Model on a custom dataset. Currently, trains a TabNet Model.
+     *
+     * <p>In order to train on a custom dataset, you must create a custom {@link TabularDataset} to
+     * load your data.
+     *
+     * @param dataset the data to train with
+     * @param performance to determine the desired model tradeoffs
+     * @return the model as a {@link ZooModel}
+     * @throws IOException if the dataset could not be loaded
+     * @throws TranslateException if the translator has errors
+     */
+    public static ZooModel<NDList, NDList> train(TabularDataset dataset, Performance performance)
+            throws IOException, TranslateException {
+        Dataset[] splitDataset = dataset.randomSplit(8, 2);
+        Dataset trainDataset = splitDataset[0];
+        Dataset validateDataset = splitDataset[1];
+        int featureSize = dataset.getFeatureSize();
+        int labelSize = dataset.getLabelSize();
+
+        Block block;
+        if (performance.equals(Performance.FAST)) {
+            // for fast cases, we set the number of independent layers and shared layers lower
+            block =
+                    TabNet.builder()
+                            .setInputDim(featureSize)
+                            .setOutDim(labelSize)
+                            .optNumIndependent(1)
+                            .optNumShared(1)
+                            .build();
+        } else if (performance.equals(Performance.BALANCED)) {
+            block = TabNet.builder().setInputDim(featureSize).setOutDim(labelSize).build();
+        } else {
+            // for accurate cases, we set the number of independent layers and shared layers higher
+            block =
+                    TabNet.builder()
+                            .setInputDim(featureSize)
+                            .setOutDim(labelSize)
+                            .optNumIndependent(4)
+                            .optNumShared(4)
+                            .build();
+        }
+
+        Model model = Model.newInstance("tabular");
+        model.setBlock(block);
+
+        TrainingConfig trainingConfig =
+                new DefaultTrainingConfig(new TabNetRegressionLoss())
+                        .addTrainingListeners(TrainingListener.Defaults.basic());
+
+        try (Trainer trainer = model.newTrainer(trainingConfig)) {
+            trainer.initialize(new Shape(1, featureSize));
+            EasyTrain.fit(trainer, 20, trainDataset, validateDataset);
+        }
+
+        return new ZooModel<>(model, new NoopTranslator());
+    }
+}
diff --git a/djl-zero/src/main/java/ai/djl/zero/tabular/package-info.java b/djl-zero/src/main/java/ai/djl/zero/tabular/package-info.java
new file mode 100644
index 000000000..42c8aebc1
--- /dev/null
+++ b/djl-zero/src/main/java/ai/djl/zero/tabular/package-info.java
@@ -0,0 +1,15 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+
+/** Contains training for Tabular({@link ai.djl.Application.Tabular}) tasks. */
+package ai.djl.zero.tabular;
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
index 5ae9e0494..192b1c444 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/JniUtils.java
@@ -1544,6 +1544,10 @@ public final class JniUtils {
             String[] extraFileKeys,
             String[] extraFileValues) {
         Device device = manager.getDevice();
+        // MPS doesn't support mapLocation
+        if ("mps".equals(device.getDeviceType())) {
+            mapLocation = false;
+        }
         logger.debug("mapLocation: {}", mapLocation);
         logger.debug("extraFileKeys: {}", Arrays.toString(extraFileKeys));
         long handle =
@@ -1571,6 +1575,11 @@ public final class JniUtils {
         if (hasSize) {
             size = new DataInputStream(is).readLong();
         }
+        // MPS doesn't support mapLocation
+        if ("mps".equals(device.getDeviceType())) {
+            mapLocation = false;
+        }
+        logger.debug("mapLocation: {}", mapLocation);
         return PyTorchLibrary.LIB.moduleLoad(
                 is,
                 new int[] {PtDeviceType.toDeviceType(device), device.getDeviceId()},
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java
index 331e1ff52..3eec14fbc 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/SavedModelBundle.java
@@ -14,17 +14,31 @@
 package ai.djl.tensorflow.engine;
 
 import org.tensorflow.internal.c_api.TF_Graph;
+import org.tensorflow.internal.c_api.TF_Operation;
 import org.tensorflow.internal.c_api.TF_Session;
+import org.tensorflow.internal.c_api.global.tensorflow;
+import org.tensorflow.proto.framework.CollectionDef;
 import org.tensorflow.proto.framework.MetaGraphDef;
+import org.tensorflow.proto.framework.SignatureDef;
 
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /** The wrapper class for native resources required for SavedModelBundle. */
 public class SavedModelBundle implements AutoCloseable {
 
+    private static final String INIT_OP_SIGNATURE_KEY = "__saved_model_init_op";
+    private static final String MAIN_OP_COLLECTION_KEY = "saved_model_main_op";
+    private static final String LEGACY_INIT_OP_COLLECTION_KEY = "legacy_init_op";
+    private static final String TABLE_INITIALIZERS_COLLECTION_KEY = "table_initializer";
+
     private TF_Graph graphHandle;
     private TF_Session sessionHandle;
     private MetaGraphDef metaGraphDef;
+    private TF_Operation[] targetOpHandles;
     private AtomicBoolean closed;
 
     public SavedModelBundle(
@@ -33,6 +47,61 @@ public class SavedModelBundle implements AutoCloseable {
         this.sessionHandle = sessionHandle;
         this.metaGraphDef = metaGraphDef;
         closed = new AtomicBoolean(false);
+
+        Map<String, SignatureDef> functions = new ConcurrentHashMap<>();
+        metaGraphDef
+                .getSignatureDefMap()
+                .forEach(
+                        (signatureName, signatureDef) -> {
+                            if (!functions.containsKey(signatureName)) {
+                                functions.put(signatureName, signatureDef);
+                            }
+                        });
+
+        List<TF_Operation> initOps = new ArrayList<>();
+        TF_Operation initOp = findInitOp(functions, metaGraphDef.getCollectionDefMap());
+        if (initOp != null) {
+            initOps.add(initOp);
+        }
+
+        if (metaGraphDef.containsCollectionDef(TABLE_INITIALIZERS_COLLECTION_KEY)) {
+            metaGraphDef
+                    .getCollectionDefMap()
+                    .get(TABLE_INITIALIZERS_COLLECTION_KEY)
+                    .getNodeList()
+                    .getValueList()
+                    .forEach(
+                            node -> {
+                                initOps.add(tensorflow.TF_GraphOperationByName(graphHandle, node));
+                            });
+        }
+        targetOpHandles = initOps.toArray(new TF_Operation[0]);
+    }
+
+    private TF_Operation findInitOp(
+            Map<String, SignatureDef> signatures, Map<String, CollectionDef> collections) {
+        SignatureDef initSig = signatures.get(INIT_OP_SIGNATURE_KEY);
+        if (initSig != null) {
+            String opName = initSig.getOutputsMap().get(INIT_OP_SIGNATURE_KEY).getName();
+            return tensorflow.TF_GraphOperationByName(graphHandle, opName);
+        }
+
+        CollectionDef initCollection;
+        if (collections.containsKey(MAIN_OP_COLLECTION_KEY)) {
+            initCollection = collections.get(MAIN_OP_COLLECTION_KEY);
+        } else {
+            initCollection = collections.get(LEGACY_INIT_OP_COLLECTION_KEY);
+        }
+
+        if (initCollection != null) {
+            CollectionDef.NodeList nodes = initCollection.getNodeList();
+            if (nodes.getValueCount() != 1) {
+                throw new IllegalArgumentException("Expected exactly one main op in saved model.");
+            }
+            String opName = nodes.getValue(0);
+            return tensorflow.TF_GraphOperationByName(graphHandle, opName);
+        }
+        return null;
     }
 
     /**
@@ -62,6 +131,10 @@ public class SavedModelBundle implements AutoCloseable {
         return metaGraphDef;
     }
 
+    TF_Operation[] getTargetOpHandles() {
+        return targetOpHandles;
+    }
+
     /** {@inheritDoc} */
     @Override
     public void close() {
diff --git a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java
index b6a853bdc..2453e58c2 100644
--- a/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java
+++ b/engines/tensorflow/tensorflow-engine/src/main/java/ai/djl/tensorflow/engine/TfSymbolBlock.java
@@ -67,6 +67,7 @@ public class TfSymbolBlock extends AbstractSymbolBlock implements AutoCloseable
         this.bundle = bundle;
         graphHandle = bundle.getGraph();
         sessionHandle = bundle.getSession();
+        targetOpHandles = bundle.getTargetOpHandles();
         MetaGraphDef metaGraphDef = bundle.getMetaGraphDef();
         Map<String, SignatureDef> signatureDefMap = metaGraphDef.getSignatureDefMap();
         if (signatureDefMap.containsKey(signatureDefKey)) {
@@ -89,8 +90,6 @@ public class TfSymbolBlock extends AbstractSymbolBlock implements AutoCloseable
         }
         describeInput();
         describeOutput();
-        // we don't use target for now
-        targetOpHandles = new TF_Operation[0];
     }
 
     /** {@inheritDoc} */
@@ -215,10 +214,6 @@ public class TfSymbolBlock extends AbstractSymbolBlock implements AutoCloseable
             for (String key : keys) {
                 TensorInfo tensorInfo = outputsMap.get(key);
                 TensorShapeProto shapeProto = tensorInfo.getTensorShape();
-                // does not support string tensors
-                if (tensorInfo.getDtype() == org.tensorflow.proto.framework.DataType.DT_STRING) {
-                    continue;
-                }
                 outputDescriptions.add(
                         key,
                         new Shape(
diff --git a/examples/src/main/java/ai/djl/examples/training/TrainAirfoilWithTabNet.java b/examples/src/main/java/ai/djl/examples/training/TrainAirfoilWithTabNet.java
new file mode 100644
index 000000000..77e3fae0c
--- /dev/null
+++ b/examples/src/main/java/ai/djl/examples/training/TrainAirfoilWithTabNet.java
@@ -0,0 +1,111 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.examples.training;
+
+import ai.djl.Model;
+import ai.djl.basicdataset.tabular.AirfoilRandomAccess;
+import ai.djl.basicmodelzoo.tabular.TabNet;
+import ai.djl.engine.Engine;
+import ai.djl.examples.training.util.Arguments;
+import ai.djl.metric.Metrics;
+import ai.djl.ndarray.types.Shape;
+import ai.djl.nn.Block;
+import ai.djl.training.DefaultTrainingConfig;
+import ai.djl.training.EasyTrain;
+import ai.djl.training.Trainer;
+import ai.djl.training.TrainingResult;
+import ai.djl.training.dataset.Dataset;
+import ai.djl.training.dataset.RandomAccessDataset;
+import ai.djl.training.listener.SaveModelTrainingListener;
+import ai.djl.training.listener.TrainingListener;
+import ai.djl.training.loss.TabNetRegressionLoss;
+import ai.djl.training.util.ProgressBar;
+import ai.djl.translate.TranslateException;
+
+import java.io.IOException;
+
+public final class TrainAirfoilWithTabNet {
+    private TrainAirfoilWithTabNet() {}
+
+    public static void main(String[] args) throws TranslateException, IOException {
+        TrainAirfoilWithTabNet.runExample(args);
+    }
+
+    public static TrainingResult runExample(String[] args) throws IOException, TranslateException {
+        Arguments arguments = new Arguments().parseArgs(args);
+        if (arguments == null) {
+            return null;
+        }
+
+        // Construct a tabNet instance
+        Block tabNet = TabNet.builder().setInputDim(5).setOutDim(1).build();
+
+        try (Model model = Model.newInstance("tabNet")) {
+            model.setBlock(tabNet);
+
+            // get the training and validation dataset
+            RandomAccessDataset[] randomAccessDatasets = getDataset(arguments);
+            RandomAccessDataset trainingSet = randomAccessDatasets[0];
+            RandomAccessDataset validateSet = randomAccessDatasets[1];
+
+            // setup training configuration
+            DefaultTrainingConfig config = setupTrainingConfig(arguments);
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                trainer.setMetrics(new Metrics());
+
+                Shape inputShape = new Shape(arguments.getBatchSize(), 5);
+
+                // initialize trainer with proper input shape
+                trainer.initialize(inputShape);
+
+                EasyTrain.fit(trainer, arguments.getEpoch(), trainingSet, validateSet);
+
+                return trainer.getTrainingResult();
+            }
+        }
+    }
+
+    private static DefaultTrainingConfig setupTrainingConfig(Arguments arguments) {
+        String outputDir = arguments.getOutputDir();
+        SaveModelTrainingListener listener = new SaveModelTrainingListener(outputDir);
+        listener.setSaveModelCallback(
+                trainer -> {
+                    TrainingResult result = trainer.getTrainingResult();
+                    Model model = trainer.getModel();
+                    model.setProperty("Loss", String.format("%.5f", result.getValidateLoss()));
+                });
+
+        return new DefaultTrainingConfig(new TabNetRegressionLoss())
+                .optDevices(Engine.getInstance().getDevices(arguments.getMaxGpus()))
+                .addTrainingListeners(TrainingListener.Defaults.logging(outputDir))
+                .addTrainingListeners(listener);
+    }
+
+    private static RandomAccessDataset[] getDataset(Arguments arguments)
+            throws IOException, TranslateException {
+        AirfoilRandomAccess.Builder airfoilBuilder = AirfoilRandomAccess.builder();
+
+        // only train dataset is available, so we get train dataset and split them
+        airfoilBuilder.optUsage(Dataset.Usage.TRAIN).setSampling(arguments.getBatchSize(), true);
+
+        for (int i = 0; i < airfoilBuilder.getAvailableFeatures().size() - 1; i++) {
+            airfoilBuilder.addFeature(airfoilBuilder.getAvailableFeatures().get(i));
+        }
+
+        AirfoilRandomAccess airfoilRandomAccess = airfoilBuilder.build();
+        airfoilRandomAccess.prepare(new ProgressBar());
+        // split the dataset into
+        return airfoilRandomAccess.randomSplit(8, 2);
+    }
+}
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/TabNet.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/TabNet.java
index dac0827ff..be2f8d007 100644
--- a/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/TabNet.java
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/TabNet.java
@@ -69,10 +69,7 @@ public final class TabNet extends AbstractBlock {
         batchNorm =
                 addChildBlock(
                         "batchNorm",
-                        BatchNorm.builder()
-                                .optAxis(0)
-                                .optMomentum(builder.batchNormMomentum)
-                                .build());
+                        BatchNorm.builder().optMomentum(builder.batchNormMomentum).build());
         List<Block> sharedBlocks = new ArrayList<>();
         for (int i = 0; i < builder.numShared; i++) {
             sharedBlocks.add(
@@ -169,7 +166,6 @@ public final class TabNet extends AbstractBlock {
         featureBlock
                 .add(
                         GhostBatchNorm.builder()
-                                .optAxis(0)
                                 .optVirtualBatchSize(virtualBatchSize)
                                 .optMomentum(batchNormMomentum)
                                 .build())
@@ -237,6 +233,7 @@ public final class TabNet extends AbstractBlock {
             PairList<String, Object> params) {
         NDManager manager = inputs.getManager();
         NDArray input = inputs.singletonOrThrow();
+        input = input.reshape(input.size(0), input.size() / input.size(0)); // batch flatten
         NDArray x =
                 batchNorm.forward(parameterStore, new NDList(input), training).singletonOrThrow();
         NDArray xa =
@@ -347,11 +344,10 @@ public final class TabNet extends AbstractBlock {
                     addChildBlock(
                             "ghostBatchNorm",
                             GhostBatchNorm.builder()
-                                    .optAxis(0)
                                     .optVirtualBatchSize(virtualBatchSize)
                                     .optMomentum(batchNormMomentum)
                                     .build());
-            sparseMax = addChildBlock("sparseMax", new SparseMax(-1, 8));
+            sparseMax = addChildBlock("sparseMax", new SparseMax());
         }
 
         /** {@inheritDoc} */
diff --git a/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/package-info.java b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/package-info.java
index ff3201da0..3046799c8 100644
--- a/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/package-info.java
+++ b/model-zoo/src/main/java/ai/djl/basicmodelzoo/tabular/package-info.java
@@ -12,7 +12,7 @@
  */
 
 /**
- * Contains classes for the computer vision section ({@link ai.djl.Application.Tabular}) of the
- * {@link ai.djl.basicmodelzoo.BasicModelZoo}.
+ * Contains classes for the tabular section ({@link ai.djl.Application.Tabular}) of the {@link
+ * ai.djl.basicmodelzoo.BasicModelZoo}.
  */
 package ai.djl.basicmodelzoo.tabular;


Test patch:
diff --git a/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java b/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java
index 649ca193b..3e8937de7 100644
--- a/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java
+++ b/engines/tensorflow/tensorflow-model-zoo/src/test/java/ai/djl/tensorflow/integration/modality/cv/TfSsdTest.java
@@ -19,11 +19,16 @@ import ai.djl.modality.Classifications;
 import ai.djl.modality.cv.Image;
 import ai.djl.modality.cv.ImageFactory;
 import ai.djl.modality.cv.output.DetectedObjects;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.NDManager;
+import ai.djl.ndarray.types.DataType;
 import ai.djl.ndarray.types.Shape;
 import ai.djl.repository.zoo.Criteria;
 import ai.djl.repository.zoo.ZooModel;
 import ai.djl.testing.TestRequirements;
 import ai.djl.training.util.ProgressBar;
+import ai.djl.translate.NoopTranslator;
 import ai.djl.translate.TranslateException;
 import ai.djl.util.Pair;
 
@@ -59,14 +64,21 @@ public class TfSsdTest {
                 Predictor<Image, DetectedObjects> predictor = model.newPredictor()) {
             Assert.assertEquals(model.describeInput().get(0).getValue(), new Shape(-1, -1, -1, 3));
             for (Pair<String, Shape> pair : model.describeOutput()) {
-                if (pair.getKey().contains("label")) {
-                    Assert.assertEquals(pair.getValue(), new Shape(-1, 1));
-                } else if (pair.getKey().contains("box")) {
-                    Assert.assertEquals(pair.getValue(), new Shape(-1, 4));
-                } else if (pair.getKey().contains("score")) {
-                    Assert.assertEquals(pair.getValue(), new Shape(-1, 1));
-                } else {
-                    throw new IllegalStateException("Unexpected output name:" + pair.getKey());
+                switch (pair.getKey()) {
+                    case "box":
+                    case "detection_boxes":
+                        Assert.assertEquals(pair.getValue(), new Shape(-1, 4));
+                        break;
+                    case "label":
+                    case "score":
+                    case "detection_class_entities":
+                    case "detection_class_labels":
+                    case "detection_class_names":
+                    case "detection_scores":
+                        Assert.assertEquals(pair.getValue(), new Shape(-1, 1));
+                        break;
+                    default:
+                        throw new IllegalStateException("Unexpected output name: " + pair.getKey());
                 }
             }
 
@@ -82,6 +94,32 @@ public class TfSsdTest {
         }
     }
 
+    @Test
+    public void testStringInputOutput() throws IOException, ModelException, TranslateException {
+        TestRequirements.notArm();
+
+        Criteria<NDList, NDList> criteria =
+                Criteria.builder()
+                        .optApplication(Application.CV.OBJECT_DETECTION)
+                        .setTypes(NDList.class, NDList.class)
+                        .optArtifactId("ssd")
+                        .optFilter("backbone", "mobilenet_v2")
+                        .optEngine("TensorFlow")
+                        .optProgress(new ProgressBar())
+                        .optTranslator(new NoopTranslator())
+                        .build();
+
+        try (ZooModel<NDList, NDList> model = criteria.loadModel();
+                Predictor<NDList, NDList> predictor = model.newPredictor();
+                NDManager manager = NDManager.newBaseManager()) {
+            NDArray array = manager.zeros(new Shape(1, 224, 224, 3));
+            NDList output = predictor.predict(new NDList(array));
+            Assert.assertEquals(output.size(), 5);
+            NDArray entities = output.get("detection_class_entities");
+            Assert.assertEquals(entities.getDataType(), DataType.STRING);
+        }
+    }
+
     private static void saveBoundingBoxImage(Image img, DetectedObjects detection)
             throws IOException {
         Path outputDir = Paths.get("build/output");
diff --git a/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java b/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
new file mode 100644
index 000000000..89a2dadcc
--- /dev/null
+++ b/examples/src/test/java/ai/djl/examples/training/TrainAirfoilWithTabNetTest.java
@@ -0,0 +1,34 @@
+/*
+ * Copyright 2022 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"). You may not use this file except in compliance
+ * with the License. A copy of the License is located at
+ *
+ * http://aws.amazon.com/apache2.0/
+ *
+ * or in the "license" file accompanying this file. This file is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES
+ * OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions
+ * and limitations under the License.
+ */
+package ai.djl.examples.training;
+
+import ai.djl.testing.TestRequirements;
+import ai.djl.training.TrainingResult;
+import ai.djl.translate.TranslateException;
+
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+import java.io.IOException;
+
+public class TrainAirfoilWithTabNetTest {
+    @Test
+    public void testTrainAirfoilWithTabNet() throws TranslateException, IOException {
+        TestRequirements.engine("MXNet", "PyTorch");
+        String[] args = new String[] {"-g", "1", "-e", "20", "-b", "32"};
+        TrainingResult result = TrainAirfoilWithTabNet.runExample(args);
+        Assert.assertNotNull(result);
+        float loss = result.getValidateLoss();
+        Assert.assertTrue(loss < 50f, "Loss: " + loss);
+    }
+}
diff --git a/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/TabNetTest.java b/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/TabNetTest.java
index 5dd6d2344..de66efd2d 100644
--- a/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/TabNetTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/model_zoo/tabular/TabNetTest.java
@@ -29,7 +29,7 @@ import ai.djl.training.TrainingConfig;
 import ai.djl.training.dataset.Batch;
 import ai.djl.training.initializer.Initializer;
 import ai.djl.training.loss.Loss;
-import ai.djl.training.loss.TabNetLoss;
+import ai.djl.training.loss.TabNetRegressionLoss;
 import ai.djl.translate.Batchifier;
 import ai.djl.util.PairList;
 
@@ -61,14 +61,14 @@ public class TabNetTest {
     @Test
     public void testTrainingAndLogic() {
         TrainingConfig config =
-                new DefaultTrainingConfig(new TabNetLoss())
+                new DefaultTrainingConfig(new TabNetRegressionLoss())
                         .optDevices(Engine.getInstance().getDevices(2));
 
         Block tabNet = TabNet.builder().setOutDim(10).build();
         try (Model model = Model.newInstance("tabNet")) {
             model.setBlock(tabNet);
             try (Trainer trainer = model.newTrainer(config)) {
-                int batchSize = 1;
+                int batchSize = 32;
                 Shape inputShape = new Shape(batchSize, 128);
                 trainer.initialize(inputShape);
                 NDManager manager = trainer.getManager();
@@ -89,7 +89,7 @@ public class TabNetTest {
                 trainer.step();
                 // the gamma of batchNorm Layer
                 Assert.assertEquals(
-                        parameters.get(0).getValue().getArray().getShape(), new Shape(1));
+                        parameters.get(0).getValue().getArray().getShape(), new Shape(128));
 
                 // weight of shared fullyConnected Block0
                 Assert.assertEquals(