| 12 | |
| 13 | |
| 14 | def lamb_update( |
| 15 | param_group, step, exp_avg, exp_avg_sq, param, grad, bias_correction, always_adapt |
| 16 | ): |
| 17 | lr = param_group["lr"] |
| 18 | weight_decay = param_group["weight_decay"] |
| 19 | eps = param_group["eps"] |
| 20 | beta0, beta1 = param_group["betas"] |
| 21 | |
| 22 | # since `conver_inputs` is disabled for param updates, |
| 23 | # scalar should be explicitly tansforred to tensor |
| 24 | |
| 25 | _lr, _neg_lr = map(tensor, (lr, -lr)) |
| 26 | _weight_decay = tensor(weight_decay) |
| 27 | _eps = tensor(eps) |
| 28 | _beta0, _beta1 = map(tensor, (beta0, beta1)) |
| 29 | |
| 30 | c1, c05, c0 = map(tensor, (1.0, 0.5, 0.0)) |
| 31 | |
| 32 | def norm(vec): |
| 33 | return sum(vec * vec) ** c05 |
| 34 | |
| 35 | p_norm = norm(param.flatten()) |
| 36 | |
| 37 | # step = step + c1 |
| 38 | step += c1 |
| 39 | |
| 40 | # exp_avg = _beta0 * exp_avg + grad * (c1 - _beta0) |
| 41 | exp_avg *= _beta0 |
| 42 | exp_avg += grad * (c1 - _beta0) |
| 43 | |
| 44 | # exp_avg_sq = _beta1 * exp_avg_sq + (c1 - _beta1) * (grad * grad) |
| 45 | exp_avg_sq *= _beta1 |
| 46 | exp_avg_sq += (c1 - _beta1) * (grad * grad) |
| 47 | |
| 48 | bias_correction1 = c1 - _beta0 ** step if bias_correction else c1 |
| 49 | bias_correction2 = c1 - _beta1 ** step if bias_correction else c1 |
| 50 | delta = (exp_avg / bias_correction1) / ( |
| 51 | (exp_avg_sq / bias_correction2) ** c05 + _eps |
| 52 | ) |
| 53 | if weight_decay != 0.0: |
| 54 | delta += param * _weight_decay |
| 55 | |
| 56 | d_norm = norm(delta.flatten()) |
| 57 | trust_ratio = ( |
| 58 | p_norm / d_norm |
| 59 | if (always_adapt or weight_decay > 0) and p_norm > c0 and d_norm > c0 |
| 60 | else c1 |
| 61 | ) |
| 62 | new_param = param - _lr * trust_ratio * delta |
| 63 | return exp_avg, exp_avg_sq, new_param |
| 64 | |
| 65 | |
| 66 | @pytest.mark.skip(reason="pytest aborted, the same as groupnorm") |