fixing base lamb optimizer

Author: Rohit Kumar Srivastava

python/mxnet/optimizer/optimizer.py

@@ -26,15 +26,15 @@
 import os
 import numpy
 from ..base import py_str
-from ..ndarray import (NDArray, zeros, clip, sqrt, cast, maximum, abs as NDabs, array, multiply,
+from ..ndarray import (NDArray, zeros, clip, sqrt, cast, minimum, maximum, abs as NDabs, array, multiply,
                        multi_sum_sq, multi_lars, norm as NDnorm)
 from ..ndarray import (sgd_update, sgd_mom_update, adam_update, rmsprop_update, rmspropalex_update,
                        mp_sgd_update, mp_sgd_mom_update, square, ftrl_update, ftml_update,
                        signsgd_update, signum_update, nag_mom_update, mp_nag_mom_update,
                        multi_sgd_update, multi_sgd_mom_update, multi_mp_sgd_update,
                        multi_mp_sgd_mom_update, preloaded_multi_sgd_update,
                        preloaded_multi_sgd_mom_update, preloaded_multi_mp_sgd_update,
-                       preloaded_multi_mp_sgd_mom_update, lamb_update)
+                       preloaded_multi_mp_sgd_mom_update, lamb_update, lamb_weight_update)
 from ..ndarray import sparse
 from ..random import normal
 from ..util import is_np_array
@@ -1250,7 +1250,7 @@ class LAMB(Optimizer):
     """LAMB Optimizer.
     """
     def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-6,
-                    lower_bound=1e-3, upper_bound=10.0, bias_correction=False, **kwargs):
+                 lower_bound=None, upper_bound=None, bias_correction=False, **kwargs):
         super(LAMB, self).__init__(learning_rate=learning_rate, **kwargs)
         self.beta1 = beta1
         self.beta2 = beta2
@@ -1259,13 +1259,14 @@ def __init__(self, learning_rate=0.001, beta1=0.9, beta2=0.999, epsilon=1e-6,
         self.upper_bound = upper_bound
         self.bias_correction = bias_correction
 
+
     def create_state(self, index, weight):
         stype = weight.stype
         dtype = weight.dtype
         return (zeros(weight.shape, weight.context, dtype=dtype, stype=stype),
                 zeros(weight.shape, weight.context, dtype=dtype, stype=stype))
 
-    def update(self, index, weight,grad, state):
+    def update(self, index, weight, grad, state):
         assert(isinstance(weight, NDArray))
         assert(isinstance(grad, NDArray))
         self._update_count(index)
@@ -1274,14 +1275,21 @@ def update(self, index, weight,grad, state):
         t = self._index_update_count[index]
 
         kwargs = {'beta1': self.beta1, 'beta2': self.beta2, 'epsilon': self.epsilon,
-                  'lower_bound': self.lower_bound, 'upper_bound': self.upper_bound,
                   'bias_correction': self.bias_correction, 't': t,
                   'rescale_grad': self.rescale_grad}
+        mean, var = state
         if self.clip_gradient:
             kwargs['clip_gradient'] = self.clip_gradient
-
-        mean, var = state
-        lamb_update(weight, grad, mean, var, out=weight, lr=lr, wd=wd, **kwargs)
+        g = lamb_update(weight, grad, mean, var, wd=wd, **kwargs)
+
+        kwargs = {}
+        if self.lower_bound:
+            kwargs['lower_bound'] = self.lower_bound
+        if self.upper_bound:
+            kwargs['upper_bound'] = self.upper_bound
+        r_1 = weight.norm()
+        r_2 = g.norm()
+        lamb_weight_update(weight, g, r_1, r_2, lr = lr, out=weight, **kwargs)
 
 
 # pylint: enable=line-too-long

src/operator/optimizer_op-inl.h

@@ -1564,20 +1564,15 @@ inline void AdamUpdateEx(const nnvm::NodeAttrs& attrs,
 }
 
 struct LAMBParam : public dmlc::Parameter<LAMBParam> {
-    float lr;
     float beta1;
     float beta2;
     float epsilon;
-    float lower_bound;
-    float upper_bound;
     float t;
     bool bias_correction;
     float wd;
     float rescale_grad;
     float clip_gradient;
     DMLC_DECLARE_PARAMETER(LAMBParam) {
-      DMLC_DECLARE_FIELD(lr)
-      .describe("Learning rate");
       DMLC_DECLARE_FIELD(beta1)
       .set_default(0.9f)
       .describe("The decay rate for the 1st moment estimates.");
@@ -1587,19 +1582,12 @@ struct LAMBParam : public dmlc::Parameter<LAMBParam> {
       DMLC_DECLARE_FIELD(epsilon)
       .set_default(1e-6f)
       .describe("A small constant for numerical stability.");
-      DMLC_DECLARE_FIELD(lower_bound)
-      .set_default(1e-3f)
-      .describe("Lower limit of norm of weight.");
-      DMLC_DECLARE_FIELD(upper_bound)
-      .set_default(10.0f)
-      .describe("Upper limit of norm of weight.");
       DMLC_DECLARE_FIELD(t)
       .describe("Index update count.");
       DMLC_DECLARE_FIELD(bias_correction)
       .set_default(false)
       .describe("Whether to use bias correction.");
       DMLC_DECLARE_FIELD(wd)
-      .set_default(0.0f)
       .describe("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.");
@@ -1614,44 +1602,48 @@ struct LAMBParam : public dmlc::Parameter<LAMBParam> {
     }
 };
 
+struct LAMBWeightParam : public dmlc::Parameter<LAMBWeightParam> {
+    float lr;
+    float lower_bound;
+    float upper_bound;
+    DMLC_DECLARE_PARAMETER(LAMBWeightParam) {
+      DMLC_DECLARE_FIELD(lr)
+      .describe("Learning rate");
+      DMLC_DECLARE_FIELD(lower_bound)
+      .set_default(1e-3f)
+      .describe("Lower limit of norm of weight.");
+      DMLC_DECLARE_FIELD(upper_bound)
+      .set_default(10.0f)
+      .describe("Upper limit of norm of weight.");
+    }
+};
+
 struct LAMBUpdateKernel {
   template<typename DType>
   MSHADOW_XINLINE static void Map(int i, DType* out_data,
     DType* mean_data, DType* var_data, const DType* weight_data, const DType* grad_data,
     const DType clip_gradient, const DType rescale_grad,
-    const DType beta1, const DType beta2,
-    DType lr, const DType wd,
-    const DType epsilon, const DType lower_bound,
-    const DType upper_bound, const DType t,
+    const DType beta1, const DType beta2,const DType wd,
+    const DType epsilon, const DType t,
     bool bias_correction, const OpReqType req) {
     using namespace mshadow_op;
 
-    DType grad_rescaled = grad_data[i] * rescale_grad + weight_data[i] * wd;
+    DType grad_rescaled = grad_data[i] * rescale_grad;
     if (clip_gradient >= 0.f) {
       grad_rescaled = clip::Map(grad_rescaled, clip_gradient);
     }
 
     mean_data[i] = beta1 * mean_data[i] + (1.f - beta1) * grad_rescaled;
     var_data[i] = beta2 * var_data[i] + (1.f - beta2) * grad_rescaled * grad_rescaled;
 
-    DType r1 = square_root::Map(square::Map(weight_data[i]));
-
-    r1 = minimum::Map(maximum::Map(r1, lower_bound), upper_bound);
-    DType g = mean_data[i] / square_root::Map(var_data[i] + epsilon) + wd * weight_data[i];
+    DType g = mean_data[i] / (square_root::Map(var_data[i]) + epsilon) + wd * weight_data[i];
 
     if (bias_correction) {
       DType mean_hat = mean_data[i] / (1. - power::Map(beta1, t));
       DType var_hat = var_data[i] / (1 - power::Map(beta2, t));
-      g = mean_hat / square_root::Map(var_hat + epsilon) + wd * weight_data[i];
-    }
-    DType r2 = square_root::Map(square::Map(g));
-    if (r1 == 0.0f || r2 == 0.0f) {
-      lr = lr * 1.0f;
-    } else {
-      lr = lr * r1 / r2;
+      g = mean_hat / (square_root::Map(var_hat) + epsilon) + wd * weight_data[i];
     }
-
-    KERNEL_ASSIGN(out_data[i], req, weight_data[i] - lr * g);
+    KERNEL_ASSIGN(out_data[i], req, g);
   }
 };
 
@@ -1661,9 +1653,9 @@ inline void LAMBUpdate(const nnvm::NodeAttrs& attrs,
                        const std::vector<TBlob> &inputs,
                        const std::vector<OpReqType> &req,
                        const std::vector<TBlob> &outputs) {
-	using namespace mxnet_op;
-	const LAMBParam& param = nnvm::get<LAMBParam>(attrs.parsed);
-	Stream<xpu>* s = ctx.get_stream<xpu>();
+  using namespace mxnet_op;
+  const LAMBParam& param = nnvm::get<LAMBParam>(attrs.parsed);
+  Stream<xpu>* s = ctx.get_stream<xpu>();
   MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
     Tensor<xpu, 2, DType> weight = inputs[0].FlatTo2D<xpu, DType>(s);
     Tensor<xpu, 2, DType> grad = inputs[1].FlatTo2D<xpu, DType>(s);
@@ -1675,13 +1667,85 @@ inline void LAMBUpdate(const nnvm::NodeAttrs& attrs,
     out.dptr_, mean.dptr_, var.dptr_, weight.dptr_, grad.dptr_,
     static_cast<DType>(param.clip_gradient), static_cast<DType>(param.rescale_grad),
     static_cast<DType>(param.beta1), static_cast<DType>(param.beta2),
-    static_cast<DType>(param.lr), static_cast<DType>(param.wd),
-    static_cast<DType>(param.epsilon), static_cast<DType>(param.lower_bound),
-    static_cast<DType>(param.upper_bound), static_cast<DType>(param.t),
-    static_cast<bool>(param.bias_correction), req[0]);
-	});
+    static_cast<DType>(param.wd), static_cast<DType>(param.epsilon),
+    static_cast<DType>(param.t), static_cast<bool>(param.bias_correction), req[0]);
+  });
+}
+
+inline bool LambWeightShape(const nnvm::NodeAttrs& attrs,
+                            mxnet::ShapeVector* in_attrs,
+                            mxnet::ShapeVector* out_attrs) {
+  CHECK_EQ(in_attrs->size(), 4U);
+  CHECK_EQ(out_attrs->size(), 1U);
+
+  mxnet::TShape expected_out(in_attrs->at(0).ndim(), -1);
+
+  mxnet::TShape& weight_shape = in_attrs->at(0);
+  mxnet::TShape& g_shape = in_attrs->at(1);
+  CHECK_EQ(weight_shape.ndim(), g_shape.ndim()) << "total no. of dimensions for weights and g must match";
+  for (int i=0; i < weight_shape.ndim(); ++i){
+    CHECK_EQ(weight_shape[i], g_shape[i]) << "weight and g dimension size mismatch at " << i << "-th index";
+  }
+  mxnet::TShape& r1_shape = in_attrs->at(2);
+  mxnet::TShape& r2_shape = in_attrs->at(3);
+  CHECK_EQ(r1_shape[0], 1U) << "r1 shape incorrect";
+  CHECK_EQ(r2_shape[0], 1U) << "r2 shape incorrect";
+  for (int i=0; i < expected_out.ndim(); ++i) {
+    expected_out[i] = weight_shape[i];
+  }
+
+  SHAPE_ASSIGN_CHECK(*out_attrs, 0, expected_out);
+  return shape_is_known(expected_out);
 }
 
+struct LAMBWeightUpdateKernel {
+  template<typename DType>
+  MSHADOW_XINLINE static void Map(int i, DType* out_data,
+    const DType* weight_data, const DType* g,
+    const DType* r1, const DType* r2,
+    DType lr, const DType lower_bound,
+    const DType upper_bound, const OpReqType req) {
+    using namespace mshadow_op;
+
+    DType new_r1 = r1[0];
+    if (lower_bound >= 0) {
+      new_r1 = maximum::Map(new_r1, lower_bound);
+    }
+    if (upper_bound >= 0) {
+      new_r1 = minimum::Map(new_r1, upper_bound);
+    }
+    if (new_r1 == 0.0f || r2[0] == 0.0f) {
+      lr = lr * 1.0f;
+    } else {
+      lr = lr * new_r1 / r2[0];
+    }
+
+    KERNEL_ASSIGN(out_data[i], req, weight_data[i] - lr * g[i]);
+  }
+};
+
+template<typename xpu>
+inline void LAMBWeightUpdate(const nnvm::NodeAttrs& attrs,
+                       const OpContext &ctx,
+                       const std::vector<TBlob> &inputs,
+                       const std::vector<OpReqType> &req,
+                       const std::vector<TBlob> &outputs) {
+  using namespace mxnet_op;
+  const LAMBWeightParam& param = nnvm::get<LAMBWeightParam>(attrs.parsed);
+  Stream<xpu>* s = ctx.get_stream<xpu>();
+  MSHADOW_REAL_TYPE_SWITCH(inputs[0].type_flag_, DType, {
+    Tensor<xpu, 2, DType> weight = inputs[0].FlatTo2D<xpu, DType>(s);
+    Tensor<xpu, 2, DType> g = inputs[1].FlatTo2D<xpu, DType>(s);
+    Tensor<xpu, 2, DType> r1 = inputs[2].FlatTo2D<xpu, DType>(s);
+    Tensor<xpu, 2, DType> r2 = inputs[3].FlatTo2D<xpu, DType>(s);
+    Tensor<xpu, 2, DType> out = outputs[0].FlatTo2D<xpu, DType>(s);
+
+  Kernel<LAMBWeightUpdateKernel, xpu>::Launch(s, weight.shape_.Size(),
+    out.dptr_, weight.dptr_, g.dptr_, r1.dptr_, r2.dptr_,
+    static_cast<DType>(param.lr), static_cast<DType>(param.lower_bound),
+    static_cast<DType>(param.upper_bound), req[0]);
+  });
+}
 
 // This RMSProp code follows the version in
 // http://arxiv.org/pdf/1308.0850v5.pdf Eq(38) - Eq(45)

src/operator/optimizer_op.cc

@@ -44,6 +44,7 @@ DMLC_REGISTER_PARAMETER(SignSGDParam);
 DMLC_REGISTER_PARAMETER(SignumParam);
 DMLC_REGISTER_PARAMETER(AdagradParam);
 DMLC_REGISTER_PARAMETER(LAMBParam);
+DMLC_REGISTER_PARAMETER(LAMBWeightParam);
 
 NNVM_REGISTER_OP(signsgd_update)
 .describe(R"code(Update function for SignSGD optimizer.
@@ -928,14 +929,33 @@ NNVM_REGISTER_OP(lamb_update)
 .set_num_inputs(4)
 .set_num_outputs(1)
 .set_attr_parser(ParamParser<LAMBParam>)
-.set_attr<mxnet::FInferShape>("FInferShape", ElemwiseShape<4,1>)
-.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<4,1>)
+.set_attr<mxnet::FInferShape>("FInferShape", ElemwiseShape<4, 1>)
+.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<4, 1>)
 .set_attr<FCompute>("FCompute<cpu>", LAMBUpdate<cpu>)
+.set_attr<nnvm::FMutateInputs>("FMutateInputs",
+  [](const nnvm::NodeAttrs& attrs) {
+    return std::vector<uint32_t>{2, 3};
+  })
 .add_argument("weight", "NDArray-or-Symbol", "Weight")
 .add_argument("grad", "NDArray-or-Symbol", "Gradient")
 .add_argument("mean", "NDArray-or-Symbol", "Moving mean")
 .add_argument("var", "NDArray-or-Symbol", "Moving variance")
 .add_arguments(LAMBParam::__FIELDS__());
 
+NNVM_REGISTER_OP(lamb_weight_update)
+.describe(R"code(Update function for lamb optimizer.
+)code" ADD_FILELINE)
+.set_num_inputs(4)
+.set_num_outputs(1)
+.set_attr_parser(ParamParser<LAMBWeightParam>)
+.set_attr<mxnet::FInferShape>("FInferShape", LambWeightShape)
+.set_attr<nnvm::FInferType>("FInferType", ElemwiseType<4, 1>)
+.set_attr<FCompute>("FCompute<cpu>", LAMBWeightUpdate<cpu>)
+.add_argument("weight", "NDArray-or-Symbol", "Weight")
+.add_argument("g", "NDArray-or-Symbol", "g")
+.add_argument("r1", "NDArray-or-Symbol", "r1")
+.add_argument("r2", "NDArray-or-Symbol", "r2")
+.add_arguments(LAMBWeightParam::__FIELDS__());
+
 }  // namespace op
 }  // namespace mxnet

src/operator/optimizer_op.cu

@@ -278,7 +278,11 @@ NNVM_REGISTER_OP(_sparse_adagrad_update)
 .set_attr<FComputeEx>("FComputeEx<gpu>", AdagradUpdateEx<gpu>);
 
 NNVM_REGISTER_OP(lamb_update)
-.set_attr<FCompute>("FCompute<gpu>", LambUpdate<gpu>);
+.set_attr<FCompute>("FCompute<gpu>", LAMBUpdate<gpu>);
+
+NNVM_REGISTER_OP(lamb_weight_update)
+.set_attr<FCompute>("FCompute<gpu>", LAMBWeightUpdate<gpu>);
+
 
 }  // namespace op
 }  // namespace mxnet

tests/python/unittest/test_optimizer.py

@@ -454,31 +454,34 @@ def update(self, index, weight, grad, state):
 
         grad *= self.rescale_grad
         if self.clip_gradient is not None:
-            grad = clip(grad, -self.clip_gradient, self.clip_gradient)
+            grad = mx.nd.clip(grad, -self.clip_gradient, self.clip_gradient)
 
         mean, var = state
         mean[:] = self.beta1 * mean + (1. - self.beta1) * grad
         var[:] = self.beta2 * var + (1. - self.beta2) * mx.nd.square(grad)
 
+        mean_hat = mean
+        var_hat = var
         r1 = weight.norm()
-        if not self.bias_correction:
-            r1 = mx.nd.minimum(mx.nd.maximum(r1, self.lower_bound), self.upper_bound)
-            g = mean / (mx.nd.sqrt(var) + self.epsilon) + wd * weight
-
-        else:
+        if self.lower_bound:
+            r1 = mx.nd.maximum(r1, self.lower_bound)
+        if self.upper_bound:
+            r1 = mx.nd.minimum(r1, self.upper_bound)
+        if self.bias_correction:
             mean_hat = mean / (1. - mx.nd.power(self.beta1, t))
             var_hat = var / (1. - mx.nd.power(self.beta2, t))
-            g = mean_hat / mx.nd.sqrt(var_hat + self.epsilon) + wd * weight
 
+        g = mean_hat / (mx.nd.sqrt(var_hat) + self.epsilon) + wd * weight
         r2 = g.norm()
-
         # calculate lamb_trust_ratio
         r = 1. if r1 == 0. or r2 == 0. else r1 / r2
         lr *= r
-
         # update weight
         weight[:] -= lr * g
 
+    def update_multi_precision(self, index, weight, grad, state):
+        self.update(index, weight, grad, state)
+
 @with_seed()
 def test_lamb():
     opt1 = PyLAMB
@@ -488,7 +491,9 @@ def test_lamb():
     rg_options = [{}, {'rescale_grad': 0.14}, {'rescale_grad': 0.8}]
     wd_options = [{}, {'wd': 0.03}, {'wd': 0.05}, {'wd': 0.07}]
     bc_options = [{}, {'bias_correction': False}, {'bias_correction': True}]
-    for params in itertools.product(cg_options, rg_options, wd_options, bc_options):
+    lb_options = [{}, {'lower_bound': None}]
+    ub_options = [{}, {'upper_bound': None}]
+    for params in itertools.product(cg_options, rg_options, wd_options, bc_options, lb_options, ub_options):
         kwarg = {k: v for param in params for k, v in param.items()}
         compare_optimizer(opt1(**kwarg), opt2(**kwarg), shape, np.float32)