Aggregated adamw update

python/mxnet/ndarray/contrib.py

@@ -548,15 +548,15 @@ def isnan(data):
     """
     return data != data  # pylint: disable=comparison-with-itself
 
-def getRescaleGrad(rescale_grad, ctx=mx.cpu()):
+def _getRescaleGrad(rescale_grad, ctx=mx.cpu()):
     if not isinstance(rescale_grad, ndarray.NDArray):
         return ndarray.full(shape=(1,), val=rescale_grad, ctx=ctx)
     else:
         return rescale_grad.as_in_context(ctx)
 
 def adamw_update(weight, grad, mean, var, rescale_grad, lr, eta, beta1=0.9, beta2=0.999,
                  epsilon=1e-8, wd=0, clip_gradient=-1, out=None, name=None, **kwargs):
-    rescale_grad = getRescaleGrad(rescale_grad, ctx=weight.context)
+    rescale_grad = _getRescaleGrad(rescale_grad, ctx=weight.context)
     return ndarray._internal._adamw_update(weight=weight, grad=grad, mean=mean, var=var,
                                            rescale_grad=rescale_grad, lr=lr, eta=eta,
                                            beta1=beta1, beta2=beta2, epsilon=epsilon,
@@ -566,7 +566,7 @@ def adamw_update(weight, grad, mean, var, rescale_grad, lr, eta, beta1=0.9, beta
 def mp_adamw_update(weight, grad, mean, var, weight32, rescale_grad, lr, eta, beta1=0.9,
                     beta2=0.999, epsilon=1e-8, wd=0, clip_gradient=-1, out=None,
                     name=None, **kwargs):
-    rescale_grad = getRescaleGrad(rescale_grad, ctx=weight.context)
+    rescale_grad = _getRescaleGrad(rescale_grad, ctx=weight.context)
     return ndarray._internal._mp_adamw_update(weight=weight, grad=grad, mean=mean, var=var,
                                               weight32=weight32,
                                               rescale_grad=rescale_grad, lr=lr, eta=eta,
@@ -579,7 +579,7 @@ def multi_adamw_update(weights, grads, mean, var, rescale_grad, lrs, wds, etas,
     if not size:
         size = len(weights)
 
-    rescale_grad = getRescaleGrad(rescale_grad, ctx=weights[0].context)
+    rescale_grad = _getRescaleGrad(rescale_grad, ctx=weights[0].context)
     temp_list = _flatten_list(zip(weights, grads, mean, var)) + [rescale_grad]
     return ndarray._internal._multi_adamw_update(*temp_list,
                                                  out=out,
@@ -595,7 +595,7 @@ def multi_mp_adamw_update(weights, grads, mean, var, weights32, rescale_grad, lr
     if not size:
         size = len(weights)
 
-    rescale_grad = getRescaleGrad(rescale_grad, ctx=weights[0].context)
+    rescale_grad = _getRescaleGrad(rescale_grad, ctx=weights[0].context)
     temp_list = _flatten_list(zip(weights, grads, mean, var, weights32)) + [rescale_grad]
     return ndarray._internal._multi_mp_adamw_update(*temp_list,
                                                     out=out,

src/operator/contrib/adamw-inl.h

@@ -78,17 +78,12 @@ inline bool MPUpdateInferShape(const nnvm::NodeAttrs& attrs,
                                mxnet::ShapeVector *out_attrs) {
   CHECK_EQ(in_attrs->size(), static_cast<size_t>(total_in)) << " in operator " << attrs.name;
   CHECK_EQ(out_attrs->size(), static_cast<size_t>(n_out)) << " in operator " << attrs.name;
-  // rescale_grad.shape = ()
   SHAPE_ASSIGN_CHECK(*in_attrs, total_in - 1, mxnet::TShape());
   // TODO(@reminisce): change "none" behavior in ElemwiseAttr
   return ElemwiseAttr<mxnet::TShape, shape_is_none, shape_assign, true, shape_string, n_in, n_out>(
       attrs, in_attrs, out_attrs, mxnet::TShape());
 }
 
-// rescale_grad is a reserved argument at position -1. Example:
-// n_in = 2: weight, grad (fp16)
-// n_out = 1: weight (fp16)
-// total_in = 6: weight, grad, mean, var, weight32, rescale_grad (fp32)
 template<int n_in, int n_out, int total_in>
 inline bool MPUpdateInferType(const nnvm::NodeAttrs& attrs,
                               std::vector<int> *in_attrs,
@@ -195,15 +190,13 @@ struct AdamWUpdate {
 ////
 // Multiple gradients in single kernel
 ////
-
 struct MultiAdamWParam : public dmlc::Parameter<MultiAdamWParam> {
   mxnet::Tuple<float> lrs;
   mxnet::Tuple<float> wds;
   mxnet::Tuple<float> etas;
   float beta1;
   float beta2;
   float epsilon;
-  float rescale_grad;
   float clip_gradient;
   int num_weights;
   DMLC_DECLARE_PARAMETER(MultiAdamWParam) {
@@ -273,7 +266,7 @@ inline bool MP_MultiAdamW_InferShape(const nnvm::NodeAttrs& attrs,
     }
     all_inferred = all_inferred && ElemwiseShape<input_stride, 1>(attrs, &input_vec, &output_vec);
   }
-  // rescale_grad.shape = ()
+
   SHAPE_ASSIGN_CHECK(*in_attrs, param.num_weights*input_stride, mxnet::TShape());
   return all_inferred;
 }
@@ -312,7 +305,7 @@ inline bool MP_MultiAdamW_InferType(const nnvm::NodeAttrs& attrs,
       TYPE_ASSIGN_CHECK(input_types, input_stride * i + input_stride - 1 - j, mshadow::kFloat32);
     }
   }
-  // rescale_grad.type = ()
+
   TYPE_ASSIGN_CHECK(input_types, param.num_weights*input_stride, mshadow::kFloat32);
   return all_inferred;
 }