【Hackathon 7th PPSCI No.12】Adam、AdamW 优化器支持 amsgrad -part

paddle/fluid/operators/fused/fused_adam_op.cc

@@ -57,6 +57,9 @@ class FusedAdamOpMaker : public framework::OpProtoAndCheckerMaker {
     AddInput("LearningRate", "(Tensor, default Tensor<float>) Learning rate");
     AddInput("Moments1", "(Tensor) Input first moments").AsDuplicable();
     AddInput("Moments2", "(Tensor) Input second moments").AsDuplicable();
+    AddInput("Moments2Max", "(Tensor) Input second moments max for amsgrad")
+        .AsDispensable()
+        .AsDuplicable();
     AddInput("Beta1Pows",
              "(Tensor, default Tensor<float>) Input beta1 power accumulator")
         .AsDuplicable();
@@ -72,6 +75,10 @@ class FusedAdamOpMaker : public framework::OpProtoAndCheckerMaker {
     AddOutput("ParamsOut", "(Tensor) Output parameters").AsDuplicable();
     AddOutput("Moments1Out", "(Tensor) Output first moments").AsDuplicable();
     AddOutput("Moments2Out", "(Tensor) Output second moments").AsDuplicable();
+    AddOutput("Moments2MaxOut",
+              "(Tensor) Output second moments max for amsgrad")
+        .AsDispensable()
+        .AsDuplicable();
     AddOutput("Beta1PowsOut", "(Tensor) Output beta1 power accumulator")
         .AsDuplicable();
     AddOutput("Beta2PowsOut", "(Tensor) Output beta2 power accumulator")
@@ -122,6 +129,10 @@ class FusedAdamOpMaker : public framework::OpProtoAndCheckerMaker {
                   "Whether to use global beta_pow for whole model instead of "
                   "creating beta_pow for each parameter.")
         .SetDefault(false);
+    AddAttr<bool>("amsgrad",
+                  "(bool, default false) "
+                  "Whether to use the AMSGrad of this algorithm.")
+        .SetDefault(false);
 
     AddComment(R"DOC(
 Adam Optimizer.

paddle/fluid/operators/ops_signature/adam_sig.cc

@@ -24,13 +24,15 @@ KernelSignature AdamOpArgumentMapping(const ArgumentMappingContext& ctx) {
                                                 "LearningRate",
                                                 "Moment1",
                                                 "Moment2",
+                                                "Moment2Max",
                                                 "Beta1Pow",
                                                 "Beta2Pow",
                                                 "MasterParam",
                                                 "SkipUpdate"};
   paddle::small_vector<const char*> out_names = {"ParamOut",
                                                  "Moment1Out",
                                                  "Moment2Out",
+                                                 "Moment2MaxOut",
                                                  "Beta1PowOut",
                                                  "Beta2PowOut",
                                                  "MasterParamOut"};
@@ -46,6 +48,7 @@ KernelSignature AdamOpArgumentMapping(const ArgumentMappingContext& ctx) {
   attr_names.emplace_back("min_row_size_to_use_multithread");
   attr_names.emplace_back("multi_precision");
   attr_names.emplace_back("use_global_beta_pow");
+  attr_names.emplace_back("amsgrad");
 
   if (ctx.IsSelectedRowsInput("Grad")) {
     return KernelSignature("adam_dense_param_sparse_grad",

paddle/fluid/operators/ops_signature/fused_adam_sig.cc

@@ -25,13 +25,15 @@ KernelSignature FusedAdamOpArgumentMapping(
                                                 "LearningRate",
                                                 "Moments1",
                                                 "Moments2",
+                                                "Moments2Max",
                                                 "Beta1Pows",
                                                 "Beta2Pows",
                                                 "MasterParams",
                                                 "SkipUpdate"};
   paddle::small_vector<const char*> out_names = {"ParamsOut",
                                                  "Moments1Out",
                                                  "Moments2Out",
+                                                 "Moments2MaxOut",
                                                  "Beta1PowsOut",
                                                  "Beta2PowsOut",
                                                  "MasterParamsOut"};
@@ -42,7 +44,8 @@ KernelSignature FusedAdamOpArgumentMapping(
                                                   "weight_decay",
                                                   "use_adamw",
                                                   "multi_precision",
-                                                  "use_global_beta_pow"};
+                                                  "use_global_beta_pow",
+                                                  "amsgrad"};
 
   return KernelSignature("fused_adam",
                          std::move(in_names),

paddle/fluid/pybind/eager_generator.cc

@@ -3344,27 +3344,31 @@ std::map<std::string, std::set<std::string>> op_passing_outs_map = {
      {"ParamOut",
       "Moment1Out",
       "Moment2Out",
+      "Moment2MaxOut",
       "Beta1PowOut",
       "Beta2PowOut",
       "MasterParamOut"}},
     {"merged_adam",
      {"ParamOut",
       "Moment1Out",
       "Moment2Out",
+      "Moment2MaxOut",
       "Beta1PowOut",
       "Beta2PowOut",
       "MasterParamOut"}},
     {"fused_adam",
      {"ParamsOut",
       "Moments1Out",
       "Moments2Out",
+      "Moments2MaxOut",
       "Beta1PowsOut",
       "Beta2PowsOut",
       "MasterParamsOut"}},
     {"adamw",
      {"ParamOut",
       "Moment1Out",
       "Moment2Out",
+      "Moment2MaxOut",
       "Beta1PowOut",
       "Beta2PowOut",
       "MasterParamOut"}},
@@ -3544,6 +3548,7 @@ std::map<std::string, std::set<std::string>> op_ins_map = {
       "LearningRate",
       "Moment1",
       "Moment2",
+      "Moment2Max",
       "Beta1Pow",
       "Beta2Pow",
       "MasterParam"}},
@@ -3553,6 +3558,7 @@ std::map<std::string, std::set<std::string>> op_ins_map = {
       "LearningRate",
       "Moment1",
       "Moment2",
+      "Moment2Max",
       "Beta1Pow",
       "Beta2Pow",
       "MasterParam"}},
@@ -3562,6 +3568,7 @@ std::map<std::string, std::set<std::string>> op_ins_map = {
       "LearningRate",
       "Moments1",
       "Moments2",
+      "Moments2Max",
       "Beta1Pows",
       "Beta2Pows",
       "MasterParams",
@@ -3572,6 +3579,7 @@ std::map<std::string, std::set<std::string>> op_ins_map = {
       "LearningRate",
       "Moment1",
       "Moment2",
+      "Moment2Max",
       "Beta1Pow",
       "Beta2Pow",
       "MasterParam"}},
@@ -3723,27 +3731,31 @@ std::map<std::string, std::set<std::string>> op_outs_map = {
      {"ParamOut",
       "Moment1Out",
       "Moment2Out",
+      "Moment2MaxOut",
       "Beta1PowOut",
       "Beta2PowOut",
       "MasterParamOut"}},
     {"merged_adam",
      {"ParamOut",
       "Moment1Out",
       "Moment2Out",
+      "Moment2MaxOut",
       "Beta1PowOut",
       "Beta2PowOut",
       "MasterParamOut"}},
     {"fused_adam",
      {"ParamsOut",
       "Moments1Out",
       "Moments2Out",
+      "Moments2MaxOut",
       "Beta1PowsOut",
       "Beta2PowsOut",
       "MasterParamsOut"}},
     {"adamw",
      {"ParamOut",
       "Moment1Out",
       "Moment2Out",
+      "Moment2MaxOut",
       "Beta1PowOut",
       "Beta2PowOut",
       "MasterParamOut"}},

paddle/phi/infermeta/multiary.cc

@@ -152,6 +152,7 @@ void AdamInferMeta(const MetaTensor& param,
                    const MetaTensor& learning_rate,
                    const MetaTensor& moment1,
                    const MetaTensor& moment2,
+                   const MetaTensor& moment2_max,
                    const MetaTensor& beta1_pow,
                    const MetaTensor& beta2_pow,
                    const MetaTensor& master_param,
@@ -163,9 +164,11 @@ void AdamInferMeta(const MetaTensor& param,
                    int64_t min_row_size_to_use_multithread,
                    bool multi_precision,
                    bool use_global_beta_pow,
+                   bool amsgrad,
                    MetaTensor* param_out,
                    MetaTensor* moment1_out,
                    MetaTensor* moment2_out,
+                   MetaTensor* moment2_max_out,
                    MetaTensor* beta1_pow_out,
                    MetaTensor* beta2_pow_out,
                    MetaTensor* master_param_outs) {
@@ -232,6 +235,10 @@ void AdamInferMeta(const MetaTensor& param,
   moment1_out->set_dtype(moment1.dtype());
   moment2_out->set_dims(param_dims);
   moment2_out->set_dtype(moment2.dtype());
+  if (amsgrad) {
+    moment2_max_out->set_dims(param_dims);
+    moment2_max_out->set_dtype(moment2.dtype());
+  }
 
   beta1_pow_out->set_dims(beta1_pow_dims);
   beta1_pow_out->set_dtype(beta1_pow.dtype());
@@ -328,6 +335,7 @@ void AdamwInferMeta(const MetaTensor& param,
                     const MetaTensor& learning_rate,
                     const MetaTensor& moment1,
                     const MetaTensor& moment2,
+                    const MetaTensor& moment2_max,
                     const MetaTensor& beta1_pow,
                     const MetaTensor& beta2_pow,
                     const MetaTensor& master_param,
@@ -342,9 +350,11 @@ void AdamwInferMeta(const MetaTensor& param,
                     int64_t min_row_size_to_use_multithread,
                     bool multi_precision,
                     bool use_global_beta_pow,
+                    bool amsgrad,
                     MetaTensor* param_out,
                     MetaTensor* moment1_out,
                     MetaTensor* moment2_out,
+                    MetaTensor* moment2_max_out,
                     MetaTensor* beta1_pow_out,
                     MetaTensor* beta2_pow_out,
                     MetaTensor* master_param_outs) {
@@ -353,6 +363,7 @@ void AdamwInferMeta(const MetaTensor& param,
                 learning_rate,
                 moment1,
                 moment2,
+                moment2_max,
                 beta1_pow,
                 beta2_pow,
                 master_param,
@@ -364,9 +375,11 @@ void AdamwInferMeta(const MetaTensor& param,
                 min_row_size_to_use_multithread,
                 multi_precision,
                 use_global_beta_pow,
+                amsgrad,
                 param_out,
                 moment1_out,
                 moment2_out,
+                moment2_max_out,
                 beta1_pow_out,
                 beta2_pow_out,
                 master_param_outs);
@@ -3856,6 +3869,7 @@ void MergedAdamInferMeta(
     const std::vector<const MetaTensor*>& learning_rate,
     const std::vector<const MetaTensor*>& moment1,
     const std::vector<const MetaTensor*>& moment2,
+    const paddle::optional<std::vector<const MetaTensor*>>& moment2_max,
     const std::vector<const MetaTensor*>& beta1_pow,
     const std::vector<const MetaTensor*>& beta2_pow,
     const paddle::optional<std::vector<const MetaTensor*>>& master_param,
@@ -3864,9 +3878,11 @@ void MergedAdamInferMeta(
     const Scalar& epsilon,
     bool multi_precision,
     bool use_global_beta_pow,
+    bool amsgrad,
     std::vector<MetaTensor*> param_out,
     std::vector<MetaTensor*> moment1_out,
     std::vector<MetaTensor*> moment2_out,
+    std::vector<MetaTensor*> moment2_max_out,
     std::vector<MetaTensor*> beta1_pow_out,
     std::vector<MetaTensor*> beta2_pow_out,
     std::vector<MetaTensor*> master_param_out) {}
@@ -5784,6 +5800,7 @@ void FusedAdamInferMeta(
     const MetaTensor& learning_rate,
     const std::vector<const MetaTensor*>& moments1,
     const std::vector<const MetaTensor*>& moments2,
+    const paddle::optional<std::vector<const MetaTensor*>>& moments2_max,
     const std::vector<const MetaTensor*>& beta1_pows,
     const std::vector<const MetaTensor*>& beta2_pows,
     const paddle::optional<std::vector<const MetaTensor*>>& master_params,
@@ -5796,9 +5813,11 @@ void FusedAdamInferMeta(
     bool use_adamw,
     bool multi_precision,
     bool use_global_beta_pow,
+    bool amsgrad,
     std::vector<MetaTensor*> params_out,
     std::vector<MetaTensor*> moments1_out,
     std::vector<MetaTensor*> moments2_out,
+    std::vector<MetaTensor*> moments2_max_out,
     std::vector<MetaTensor*> beta1_pows_out,
     std::vector<MetaTensor*> beta2_pows_out,
     std::vector<MetaTensor*> master_params_out) {
@@ -5810,6 +5829,10 @@ void FusedAdamInferMeta(
     moments1_out[i]->set_dtype(moments1[i]->dtype());
     moments2_out[i]->set_dims(moments2[i]->dims());
     moments2_out[i]->set_dtype(moments2[i]->dtype());
+    if (amsgrad) {
+      moments2_max_out[i]->set_dims(moments2_max.get()[i]->dims());
+      moments2_max_out[i]->set_dtype(moments2_max.get()[i]->dtype());
+    }
     beta1_pows_out[i]->set_dims(beta1_pows[i]->dims());
     beta1_pows_out[i]->set_dtype(beta1_pows[i]->dtype());
     beta2_pows_out[i]->set_dims(beta2_pows[i]->dims());

paddle/phi/infermeta/multiary.h

@@ -86,6 +86,7 @@ void AdamInferMeta(const MetaTensor& param,
                    const MetaTensor& learning_rate,
                    const MetaTensor& moment1,
                    const MetaTensor& moment2,
+                   const MetaTensor& moment2_max,
                    const MetaTensor& beta1_pow,
                    const MetaTensor& beta2_pow,
                    const MetaTensor& master_param,
@@ -97,9 +98,11 @@ void AdamInferMeta(const MetaTensor& param,
                    int64_t min_row_size_to_use_multithread,
                    bool multi_precision,
                    bool use_global_beta_pow,
+                   bool amsgrad,
                    MetaTensor* param_out,
                    MetaTensor* moment1_out,
                    MetaTensor* moment2_out,
+                   MetaTensor* moment2_max_out,
                    MetaTensor* beta1_pow_out,
                    MetaTensor* beta2_pow_out,
                    MetaTensor* master_param_outs);
@@ -109,6 +112,7 @@ void AdamwInferMeta(const MetaTensor& param,
                     const MetaTensor& learning_rate,
                     const MetaTensor& moment1,
                     const MetaTensor& moment2,
+                    const MetaTensor& moment2_max,
                     const MetaTensor& beta1_pow,
                     const MetaTensor& beta2_pow,
                     const MetaTensor& master_param,
@@ -123,9 +127,11 @@ void AdamwInferMeta(const MetaTensor& param,
                     int64_t min_row_size_to_use_multithread,
                     bool multi_precision,
                     bool use_global_beta_pow,
+                    bool amsgrad,
                     MetaTensor* param_out,
                     MetaTensor* moment1_out,
                     MetaTensor* moment2_out,
+                    MetaTensor* moment2_max_out,
                     MetaTensor* beta1_pow_out,
                     MetaTensor* beta2_pow_out,
                     MetaTensor* master_param_outs);
@@ -711,6 +717,7 @@ void MergedAdamInferMeta(
     const std::vector<const MetaTensor*>& learning_rate,
     const std::vector<const MetaTensor*>& moment1,
     const std::vector<const MetaTensor*>& moment2,
+    const paddle::optional<std::vector<const MetaTensor*>>& moment2_max,
     const std::vector<const MetaTensor*>& beta1_pow,
     const std::vector<const MetaTensor*>& beta2_pow,
     const paddle::optional<std::vector<const MetaTensor*>>& master_param,
@@ -719,9 +726,11 @@ void MergedAdamInferMeta(
     const Scalar& epsilon,
     bool multi_precision,
     bool use_global_beta_pow,
+    bool amsgrad,
     std::vector<MetaTensor*> param_out,
     std::vector<MetaTensor*> moment1_out,
     std::vector<MetaTensor*> moment2_out,
+    std::vector<MetaTensor*> moment2_max_out,
     std::vector<MetaTensor*> beta1_pow_out,
     std::vector<MetaTensor*> beta2_pow_out,
     std::vector<MetaTensor*> master_param_out);
@@ -1117,6 +1126,7 @@ void FusedAdamInferMeta(
     const MetaTensor& learning_rate,
     const std::vector<const MetaTensor*>& moments1,
     const std::vector<const MetaTensor*>& moments2,
+    const paddle::optional<std::vector<const MetaTensor*>>& moments2_max,
     const std::vector<const MetaTensor*>& beta1_pows,
     const std::vector<const MetaTensor*>& beta2_pows,
     const paddle::optional<std::vector<const MetaTensor*>>& master_params,
@@ -1129,9 +1139,11 @@ void FusedAdamInferMeta(
     bool use_adamw,
     bool multi_precision,
     bool use_global_beta_pow,
+    bool amsgrad,
     std::vector<MetaTensor*> params_out,
     std::vector<MetaTensor*> moments1_out,
     std::vector<MetaTensor*> moments2_out,
+    std::vector<MetaTensor*> moments2_max_out,
     std::vector<MetaTensor*> beta1_pows_out,
     std::vector<MetaTensor*> beta2_pows_out,
     std::vector<MetaTensor*> master_params_out);