eric-haibin-lin commented on a change in pull request #16893: Multi-tensor LAMB
URL: https://github.com/apache/incubator-mxnet/pull/16893#discussion_r352879674
##########
File path: python/mxnet/optimizer/optimizer.py
##########
@@ -1051,6 +1052,95 @@ def update_multi_precision(self, index, weight, grad,
state):
self._update_impl(index, weight, grad, state,
multi_precision=use_multi_precision)
+@register
+class MultiLAMB(Optimizer):
+ """multiLAMB 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):
+ super(MultiLAMB, self).__init__(learning_rate=learning_rate, **kwargs)
+ self.beta1 = beta1
+ self.beta2 = beta2
+ self.epsilon = epsilon
+ self.lower_bound = lower_bound
+ self.upper_bound = upper_bound
+ self.bias_correction = bias_correction
+ self.aggregate_num = max(1, min(50,
int(os.getenv('MXNET_OPTIMIZER_AGGREGATION_SIZE', "50"))))
+
+ def create_state(self, index, weight):
+ stype = weight.stype
+ dtype = weight.dtype
+ return (zeros(weight.shape, weight.context, dtype=dtype, stype=stype),
# mean
+ zeros(weight.shape, weight.context, dtype=dtype, stype=stype),
# variance
+ zeros(weight.shape, weight.context, dtype=dtype, stype=stype))
# temp_g
Review comment:
I am a bit confused for the use of temp_g here. Is this mainly used as a
temporarly workspace? Is the value of temp_g at step `t-1` used by the update
at step `t`? If so would it be possible to allocate g in the op via resource
request? Otherwise the state size is increased by 50% when saving the training
state. On the other hand, it is inconsistent with the existing LAMB optimizer
which has (mean, var) as the state.
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
