Repository: deepjavalibrary/djl
Pull request #2206 (2022-12-02 05:06:01)

Implement AdamW on Pytorch and MXNet
## Description ##
https://github.com/deepjavalibrary/djl/issues/2185

It is currently implemented only in PyTorch both MXNet.

Old Description:
It looks like the JNA is not available on MXNet. There is a JNA called `_adamw_update`. But the number of inputs is 5 different than adam or pytorch adamw. The required parameters are different (it requires `eta` ). It needs further investigation on this JNA.

```java
WARNING: JNA: Callback ai.djl.mxnet.engine.MxParameterServer$OptimizerCallback@4f936da8 threw the following exception
ai.djl.engine.EngineException: MXNet engine call failed: MXNetError: Cannot find argument 'rescale_grad', Possible Arguments:
----------------
lr : float, required
    Learning rate
beta1 : float, optional, default=0.899999976
    The decay rate for the 1st moment estimates.
beta2 : float, optional, default=0.999000013
    The decay rate for the 2nd moment estimates.
epsilon : float, optional, default=9.99999994e-09
    A small constant for numerical stability.
wd : float, optional, default=0
    Weight decay augments the objective function with a regularization term that penalizes large weights. The penalty scales with the square of the magnitude of each weight.
eta : float, required
    Learning rate schedule multiplier
clip_gradient : float, optional, default=-1
    Clip gradient to the range of [-clip_gradient, clip_gradient] If clip_gradient <= 0, gradient clipping is turned off. grad = max(min(grad, clip_gradient), -clip_gradient).
, in operator _adamw_update(name="", beta2="0.999", rescale_grad="1.0", epsilon="1.0E-8", wd="0.01", beta1="0.9", clip_gradient="-1.0", lr="0.031622566")
```

codecov-commenter: # [Codecov](https://codecov.io/gh/deepjavalibrary/djl/pull/2206?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: **71.40**% // Decreases project coverage by **`-0.68%`** :warning:
> Coverage data is based on head [(`fbd02ab`)](https://codecov.io/gh/deepjavalibrary/djl/pull/2206?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: 72.57% of modified lines in pull request are covered.

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


```diff
@@             Coverage Diff              @@
##             master    #2206      +/-   ##
============================================
- Coverage     72.08%   71.40%   -0.69%     
- Complexity     5126     6345    +1219     
============================================
  Files           473      629     +156     
  Lines         21970    28066    +6096     
  Branches       2351     3023     +672     
============================================
+ Hits          15838    20041    +4203     
- Misses         4925     6552    +1627     
- Partials       1207     1473     +266     
```


| [Impacted Files](https://codecov.io/gh/deepjavalibrary/djl/pull/2206?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/2206/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/2206/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/2206/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/2206/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/2206/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/2206/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: |
| [.../modality/cv/translator/ImageFeatureExtractor.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2206/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvSW1hZ2VGZWF0dXJlRXh0cmFjdG9yLmphdmE=) | `0.00% <0.00%> (ø)` | |
| [.../ai/djl/modality/cv/translator/YoloTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2206/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3IvWW9sb1RyYW5zbGF0b3IuamF2YQ==) | `27.77% <0.00%> (+18.95%)` | :arrow_up: |
| [...modality/cv/translator/wrapper/FileTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2206/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3Ivd3JhcHBlci9GaWxlVHJhbnNsYXRvci5qYXZh) | `44.44% <ø> (ø)` | |
| [...y/cv/translator/wrapper/InputStreamTranslator.java](https://codecov.io/gh/deepjavalibrary/djl/pull/2206/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=deepjavalibrary#diff-YXBpL3NyYy9tYWluL2phdmEvYWkvZGpsL21vZGFsaXR5L2N2L3RyYW5zbGF0b3Ivd3JhcHBlci9JbnB1dFN0cmVhbVRyYW5zbGF0b3IuamF2YQ==) | `44.44% <ø> (ø)` | |
| ... and [565 more](https://codecov.io/gh/deepjavalibrary/djl/pull/2206/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/2206?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/ndarray/internal/NDArrayEx.java b/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
index 8368a4716..b12ac5dd0 100644
--- a/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
+++ b/api/src/main/java/ai/djl/ndarray/internal/NDArrayEx.java
@@ -227,13 +227,15 @@ public interface NDArrayEx {
             NDList inputs,
             NDList weights,
             float learningRate,
+            float learningRateBiasCorrection,
             float weightDecay,
             float rescaleGrad,
             float clipGrad,
             float beta1,
             float beta2,
             float epsilon,
-            boolean lazyUpdate);
+            boolean lazyUpdate,
+            boolean adamw);
 
     void nagUpdate(
             NDList inputs,
diff --git a/api/src/main/java/ai/djl/nn/AbstractBaseBlock.java b/api/src/main/java/ai/djl/nn/AbstractBaseBlock.java
index d20d6ba43..cb0d6ebe5 100644
--- a/api/src/main/java/ai/djl/nn/AbstractBaseBlock.java
+++ b/api/src/main/java/ai/djl/nn/AbstractBaseBlock.java
@@ -46,7 +46,7 @@ public abstract class AbstractBaseBlock implements Block {
     /** The shape of the input for this block, set by the initialization process. */
     protected Shape[] inputShapes;
 
-    protected DataType[] outputDatatypes;
+    protected DataType[] outputDataTypes;
 
     /** List of names for the input, named inputs should be manually set in sub class. */
     protected List<String> inputNames = Collections.emptyList();
@@ -375,6 +375,6 @@ public abstract class AbstractBaseBlock implements Block {
     /** {@inheritDoc} */
     @Override
     public DataType[] getOutputDataTypes() {
-        return outputDatatypes;
+        return outputDataTypes;
     }
 }
diff --git a/api/src/main/java/ai/djl/nn/LambdaBlock.java b/api/src/main/java/ai/djl/nn/LambdaBlock.java
index 26b3f38a3..c0c926499 100644
--- a/api/src/main/java/ai/djl/nn/LambdaBlock.java
+++ b/api/src/main/java/ai/djl/nn/LambdaBlock.java
@@ -122,7 +122,7 @@ public class LambdaBlock extends AbstractBlock {
                 outputShapes[i] = output.get(i).getShape();
                 dataTypes[i] = output.get(i).getDataType();
             }
-            outputDatatypes = dataTypes;
+            outputDataTypes = dataTypes;
             return outputShapes;
         }
     }
diff --git a/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java b/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
index 6e15103ca..cb5e68921 100644
--- a/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
+++ b/api/src/main/java/ai/djl/repository/zoo/BaseModelLoader.java
@@ -28,6 +28,7 @@ import ai.djl.translate.TranslatorFactory;
 import ai.djl.util.ClassLoaderUtils;
 import ai.djl.util.Pair;
 import ai.djl.util.Progress;
+import ai.djl.util.Utils;
 
 import java.io.IOException;
 import java.io.Reader;
@@ -102,6 +103,7 @@ public class BaseModelLoader implements ModelLoader {
             if (modelDir == null) {
                 throw new AssertionError("Directory should not be null.");
             }
+            modelDir = Utils.getNestedModelDir(modelDir);
 
             loadServingProperties(modelDir, arguments, options);
             Application application = criteria.getApplication();
diff --git a/api/src/main/java/ai/djl/training/optimizer/Adam.java b/api/src/main/java/ai/djl/training/optimizer/Adam.java
index 5232dbb50..22ee0b390 100644
--- a/api/src/main/java/ai/djl/training/optimizer/Adam.java
+++ b/api/src/main/java/ai/djl/training/optimizer/Adam.java
@@ -24,16 +24,15 @@ import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 
 /**
- * {@code Adam} is a generalization of the AdaGrad {@link Optimizer}.
- *
- * <p>Adam updates the weights using:<br>
- * <br>
+ * {@code Adam} is a generalization of the AdaGrad {@link Optimizer}. <br>
+ * \( grad += weight_decay * w\)<br>
  * \( m = beta1 * m + (1 - beta1) * grad\)<br>
  * \( v = beta2 * v + (1 - beta2) * grad^2 \)<br>
- * \( w -= learning_rate * m / (sqrt(v) + epsilon) \)<br>
+ * \( learning_rate_bias_correction = learning_rate / beta1**t * sqrt(beta2**t) \)<br>
+ * \( w -= learning_rate_bias_correction * m / (sqrt(v) + epsilon) \)<br>
  * <br>
  * where g represents the gradient, and m/v are 1st and 2nd order moment estimates (mean and
- * variance).
+ * variance), t is the step.
  *
  * @see <a href="https://d2l.djl.ai/chapter_optimization/adam.html">The D2L chapter on Adam</a>
  */
@@ -68,8 +67,8 @@ public class Adam extends Optimizer {
         int t = updateCount(parameterId);
         double coef1 = 1.0 - Math.pow(beta1, t);
         double coef2 = 1.0 - Math.pow(beta2, t);
-        float lr = learningRateTracker.getNewValue(parameterId, t);
-        float newLearningRate = (float) (lr * Math.sqrt(coef2) / coef1);
+        float newLearningRate = learningRateTracker.getNewValue(parameterId, t);
+        float learningRateBiasCorrection = (float) (newLearningRate * Math.sqrt(coef2) / coef1);
         float weightDecay = getWeightDecay();
 
         Preconditions.checkArgument(
@@ -97,13 +96,15 @@ public class Adam extends Optimizer {
                 inputs,
                 weights,
                 newLearningRate,
+                learningRateBiasCorrection,
                 weightDecay,
                 rescaleGrad,
                 clipGrad,
                 beta1,
                 beta2,
                 epsilon,
-                true);
+                true,
+                false);
     }
 
     /**
diff --git a/api/src/main/java/ai/djl/training/optimizer/AdamW.java b/api/src/main/java/ai/djl/training/optimizer/AdamW.java
new file mode 100644
index 000000000..f700315f0
--- /dev/null
+++ b/api/src/main/java/ai/djl/training/optimizer/AdamW.java
@@ -0,0 +1,194 @@
+/*
+ * 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.optimizer;
+
+import ai.djl.Device;
+import ai.djl.ndarray.NDArray;
+import ai.djl.ndarray.NDList;
+import ai.djl.ndarray.internal.NDArrayEx;
+import ai.djl.training.tracker.ParameterTracker;
+import ai.djl.training.tracker.Tracker;
+import ai.djl.util.Preconditions;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+/**
+ * {@code Adam} is a generalization of the AdaGrad {@link Optimizer}.
+ *
+ * <p>Adam updates the weights using:<br>
+ * <br>
+ * \( w *= (1 - learning_rate * weight_decay\)<br>
+ * \( m = beta1 * m + (1 - beta1) * grad\)<br>
+ * \( v = beta2 * v + (1 - beta2) * grad^2 \)<br>
+ * \( learning_rate_bias_correction = learning_rate / beta1**t * sqrt(beta2**t) \)<br>
+ * \( w -= learning_rate_bias_correction * m / (sqrt(v) + epsilon) \)<br>
+ * <br>
+ * where g represents the gradient, and m/v are 1st and 2nd order moment estimates (mean and
+ * variance), t is the step.
+ *
+ * @see <a href="https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html">The algorithm of
+ *     AdamW</a>
+ */
+public class AdamW extends Optimizer {
+
+    private ParameterTracker learningRateTracker;
+    private float beta1;
+    private float beta2;
+    private float epsilon;
+
+    private Map<String, Map<Device, NDArray>> means;
+    private Map<String, Map<Device, NDArray>> variances;
+
+    /**
+     * Creates a new instance of {@code Adam} optimizer.
+     *
+     * @param builder the builder to create a new instance of {@code Adam} optimizer
+     */
+    protected AdamW(Builder builder) {
+        super(builder);
+        learningRateTracker = builder.learningRateTracker;
+        beta1 = builder.beta1;
+        beta2 = builder.beta2;
+        epsilon = builder.epsilon;
+        means = new ConcurrentHashMap<>();
+        variances = new ConcurrentHashMap<>();
+    }
+
+    /** {@inheritDoc} */
+    @Override
+    public void update(String parameterId, NDArray weight, NDArray grad) {
+        int t = updateCount(parameterId);
+        double coef1 = 1.0 - Math.pow(beta1, t);
+        double coef2 = 1.0 - Math.pow(beta2, t);
+        float newLearningRate = learningRateTracker.getNewValue(parameterId, t);
+        float learningRateBiasCorrection = (float) (newLearningRate * Math.sqrt(coef2) / coef1);
+        float weightDecay = getWeightDecay();
+
+        Preconditions.checkArgument(
+                !Float.isNaN(newLearningRate)
+                        && !Float.isNaN(weightDecay)
+                        && !Float.isInfinite(newLearningRate)
+                        && !Float.isInfinite(weightDecay),
+                "learning rate or weight decay is nan or infinite");
+        NDList inputs =
+                new NDList(
+                        weight,
+                        grad,
+                        withDefaultState(
+                                means, parameterId, weight.getDevice(), k -> weight.zerosLike()),
+                        withDefaultState(
+                                variances,
+                                parameterId,
+                                weight.getDevice(),
+                                k -> weight.zerosLike()));
+        NDList weights = new NDList(weight);
+
+        NDArrayEx ex = weight.getNDArrayInternal();
+
+        ex.adamUpdate(
+                inputs,
+                weights,
+                newLearningRate,
+                learningRateBiasCorrection,
+                weightDecay,
+                rescaleGrad,
+                clipGrad,
+                beta1,
+                beta2,
+                epsilon,
+                true,
+                true);
+    }
+
+    /**
+     * Creates a builder to build a {@code Adam}.
+     *
+     * @return a new builder
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** The Builder to construct an {@link AdamW} object. */
+    public static final class Builder extends OptimizerBuilder<Builder> {
+
+        private ParameterTracker learningRateTracker = Tracker.fixed(0.001f);
+        private float beta1 = 0.9f;
+        private float beta2 = 0.999f;
+        private float epsilon = 1e-8f;
+
+        Builder() {
+            optWeightDecays(0.01f);
+        }
+
+        /** {@inheritDoc} */
+        @Override
+        protected Builder self() {
+            return this;
+        }
+
+        /**
+         * Sets the {@link ParameterTracker} for this optimizer.
+         *
+         * @param learningRateTracker the {@link ParameterTracker} to be set
+         * @return this {@code Builder}
+         */
+        public Builder optLearningRateTracker(ParameterTracker learningRateTracker) {
+            this.learningRateTracker = learningRateTracker;
+            return this;
+        }
+
+        /**
+         * Sets the decay rate for the first moment estimates.
+         *
+         * @param beta1 the deacay rate for the the first moment estimates
+         * @return this {@code Builder}
+         */
+        public Builder optBeta1(float beta1) {
+            this.beta1 = beta1;
+            return this;
+        }
+
+        /**
+         * Sets the decay rate for the second moment estimates.
+         *
+         * @param beta2 the decay rate for the the second moment estimates
+         * @return this {@code Builder}
+         */
+        public Builder optBeta2(float beta2) {
+            this.beta2 = beta2;
+            return this;
+        }
+
+        /**
+         * Sets \(epsilon\) - a small quantity for numerical stability.
+         *
+         * @param epsilon a small quantity for numerical stability
+         * @return this {@code Builder}
+         */
+        public Builder optEpsilon(float epsilon) {
+            this.epsilon = epsilon;
+            return this;
+        }
+
+        /**
+         * Builds a {@link AdamW} block.
+         *
+         * @return the {@link AdamW} block
+         */
+        public AdamW build() {
+            return new AdamW(this);
+        }
+    }
+}
diff --git a/api/src/main/java/ai/djl/training/optimizer/Optimizer.java b/api/src/main/java/ai/djl/training/optimizer/Optimizer.java
index ce3853da0..36400fa0e 100644
--- a/api/src/main/java/ai/djl/training/optimizer/Optimizer.java
+++ b/api/src/main/java/ai/djl/training/optimizer/Optimizer.java
@@ -77,6 +77,16 @@ public abstract class Optimizer {
         return new Adam.Builder();
     }
 
+    /**
+     * Returns a new instance of {@link ai.djl.training.optimizer.AdamW.Builder} that can build an
+     * {@link AdamW} optimizer.
+     *
+     * @return the {@link AdamW} {@link ai.djl.training.optimizer.AdamW.Builder}
+     */
+    public static AdamW.Builder adamW() {
+        return new AdamW.Builder();
+    }
+
     /**
      * Returns a new instance of {@link RmsProp.Builder} that can build an {@link RmsProp}
      * optimizer.
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 9dfd83110..b3d85dff5 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
@@ -424,25 +424,36 @@ class MxNDArrayEx implements NDArrayEx {
             NDList inputs,
             NDList weights,
             float learningRate,
+            float learningRateBiasCorrection,
             float weightDecay,
             float rescaleGrad,
             float clipGrad,
             float beta1,
             float beta2,
             float epsilon,
-            boolean lazyUpdate) {
+            boolean lazyUpdate,
+            boolean adamw) {
         MxOpParams params = new MxOpParams();
-        params.addParam("lr", learningRate);
-        params.addParam("wd", weightDecay);
-        params.addParam("rescale_grad", rescaleGrad);
+        params.addParam("lr", learningRateBiasCorrection);
         params.addParam("clip_gradient", clipGrad);
 
         params.addParam("beta1", beta1);
         params.addParam("beta2", beta2);
         params.addParam("epsilon", epsilon);
-        params.addParam("lazy_update", lazyUpdate);
 
-        getManager().invoke("adam_update", inputs, weights, params);
+        if (!adamw) {
+            params.addParam("wd", weightDecay);
+            params.addParam("rescale_grad", rescaleGrad);
+            params.addParam("lazy_update", lazyUpdate);
+            getManager().invoke("adam_update", inputs, weights, params);
+        } else {
+            // https://github.com/apache/mxnet/blob/7d602e3b2382eb501fdeb94c4d97e652a723af11/src/operator/contrib/adamw.cc#L80-L121
+            // https://github.com/apache/mxnet/blob/7d602e3b2382eb501fdeb94c4d97e652a723af11/src/operator/contrib/adamw-inl.h#L172-L207
+            inputs.add(inputs.getManager().create(rescaleGrad));
+            params.addParam("eta", 1.0f);
+            params.addParam("wd", weightDecay * learningRate);
+            getManager().invoke("_adamw_update", inputs, weights, params);
+        }
     }
 
     /** {@inheritDoc} */
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 28375ef38..f5ae2cdbd 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
@@ -266,13 +266,15 @@ public class PtNDArrayEx implements NDArrayEx {
             NDList inputs,
             NDList weights,
             float learningRate,
+            float learningRateBiasCorrection,
             float weightDecay,
             float rescaleGrad,
             float clipGrad,
             float beta1,
             float beta2,
             float epsilon,
-            boolean lazyUpdate) {
+            boolean lazyUpdate,
+            boolean adamw) {
         // TODO: Lazy update not used
         PtNDManager manager = array.getManager();
         JniUtils.adamUpdate(
@@ -281,12 +283,14 @@ public class PtNDArrayEx implements NDArrayEx {
                 manager.from(inputs.get(2)),
                 manager.from(inputs.get(3)),
                 learningRate,
+                learningRateBiasCorrection,
                 weightDecay,
                 rescaleGrad,
                 clipGrad,
                 beta1,
                 beta2,
-                epsilon);
+                epsilon,
+                adamw);
         // call zero-grad
         JniUtils.zeroGrad(manager.from(weights.singletonOrThrow()));
     }
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 862833d44..47fb0b724 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
@@ -1647,24 +1647,28 @@ public final class JniUtils {
             PtNDArray mean,
             PtNDArray variance,
             float lr,
+            float learningRateBiasCorrection,
             float wd,
             float rescaleGrad,
             float clipGrad,
             float beta1,
             float beta2,
-            float eps) {
+            float eps,
+            boolean adamw) {
         PyTorchLibrary.LIB.adamUpdate(
                 weight.getHandle(),
                 grad.getHandle(),
                 mean.getHandle(),
                 variance.getHandle(),
                 lr,
+                learningRateBiasCorrection,
                 wd,
                 rescaleGrad,
                 clipGrad,
                 beta1,
                 beta2,
-                eps);
+                eps,
+                adamw);
     }
 
     public static void sgdUpdate(
diff --git a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
index 381712390..53b043b4e 100644
--- a/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
+++ b/engines/pytorch/pytorch-engine/src/main/java/ai/djl/pytorch/jni/PyTorchLibrary.java
@@ -599,12 +599,14 @@ final class PyTorchLibrary {
             long mean,
             long variance,
             float lr,
+            float learningRateBiasCorrection,
             float wd,
             float rescaleGrad,
             float clipGrad,
             float beta1,
             float beta2,
-            float eps);
+            float eps,
+            boolean adamw);
 
     native void sgdUpdate(
             long weight,
diff --git a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_optim.cc b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_optim.cc
index 0b3a49358..70759b7e1 100644
--- a/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_optim.cc
+++ b/engines/pytorch/pytorch-native/src/main/native/ai_djl_pytorch_jni_PyTorchLibrary_optim.cc
@@ -18,8 +18,9 @@
 // The file is the implementation for PyTorch training operations
 
 JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_adamUpdate(JNIEnv* env, jobject jthis, jlong jweight,
-    jlong jgrad, jlong jmean, jlong jvariance, jfloat learning_rate, jfloat weight_decay, jfloat rescale_grad,
-    jfloat clip_grad, jfloat beta1, jfloat beta2, jfloat eps) {
+    jlong jgrad, jlong jmean, jlong jvariance, jfloat learning_rate, jfloat learning_rate_bias_correction,
+    jfloat weight_decay, jfloat rescale_grad, jfloat clip_grad, jfloat beta1, jfloat beta2, jfloat eps,
+    jboolean adamw) {
   API_BEGIN()
   torch::autograd::AutoGradMode no_autograd_guard{false};
   const auto* weight_ptr = reinterpret_cast<torch::Tensor*>(jweight);
@@ -34,10 +35,15 @@ JNIEXPORT void JNICALL Java_ai_djl_pytorch_jni_PyTorchLibrary_adamUpdate(JNIEnv*
     // Add clip grad option
     grad.clamp_max_(clip_grad);
   }
-  grad.add_(*weight_ptr, weight_decay);
+  if (!adamw) {
+    // rescaled_grad is obtained here
+    grad.add_(*weight_ptr, weight_decay);
+  } else {
+    weight_ptr->sub_(weight_ptr->mul(learning_rate).mul(weight_decay));
+  }
   mean_ptr->mul_(beta1).add_(grad, 1 - beta1);
   variance_ptr->mul_(beta2).addcmul_(grad, grad, 1 - beta2);
-  weight_ptr->sub_(mean_ptr->mul(learning_rate).div(variance_ptr->sqrt().add(eps)));
+  weight_ptr->sub_(mean_ptr->mul(learning_rate_bias_correction).div(variance_ptr->sqrt().add(eps)));
   API_END()
 }
 
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 8c2091137..b7c9ec3e0 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
@@ -286,13 +286,15 @@ public class TfNDArrayEx implements NDArrayEx {
             NDList inputs,
             NDList weights,
             float learningRate,
+            float learningRateBiasCorrection,
             float weightDecay,
             float rescaleGrad,
             float clipGrad,
             float beta1,
             float beta2,
             float epsilon,
-            boolean lazyUpdate) {
+            boolean lazyUpdate,
+            boolean adamw) {
         throw new UnsupportedOperationException("Not implemented");
     }
 


Test patch:
diff --git a/integration/src/main/java/ai/djl/integration/tests/training/OptimizerTest.java b/integration/src/main/java/ai/djl/integration/tests/training/OptimizerTest.java
index bc1b07470..b54225dfb 100644
--- a/integration/src/main/java/ai/djl/integration/tests/training/OptimizerTest.java
+++ b/integration/src/main/java/ai/djl/integration/tests/training/OptimizerTest.java
@@ -172,6 +172,35 @@ public class OptimizerTest {
         }
     }
 
+    @Test
+    public void testAdamW() {
+        Optimizer optim = Optimizer.adamW().optLearningRateTracker(Tracker.fixed(0.1f)).build();
+
+        Device[] devices = Engine.getInstance().getDevices(1);
+        TrainingConfig config =
+                new DefaultTrainingConfig(Loss.l2Loss())
+                        .optInitializer(Initializer.ONES, Parameter.Type.WEIGHT)
+                        .optOptimizer(optim)
+                        .optDevices(devices);
+        Block block = Linear.builder().setUnits(CHANNELS).build();
+        try (Model model = Model.newInstance("model", devices[0])) {
+            model.setBlock(block);
+
+            try (Trainer trainer = model.newTrainer(config)) {
+                int batchSize = config.getDevices().length * BATCH_SIZE;
+                trainer.initialize(new Shape(batchSize, CHANNELS));
+
+                NDManager manager = trainer.getManager();
+                NDArray result = runOptimizer(manager, trainer, block, batchSize);
+                NDArray result2 = runOptimizer(manager, trainer, block, batchSize);
+                Assertions.assertAlmostEquals(
+                        result, manager.create(new float[] {0.8990001082420349f, -0.10000064f}));
+                Assertions.assertAlmostEquals(
+                        result2, manager.create(new float[] {0.7987207174301147f, -0.19939029f}));
+            }
+        }
+    }
+
     @Test
     public void testAdagrad() {
         Optimizer optim = Optimizer.adagrad().optLearningRateTracker(Tracker.fixed(0.1f)).build();