| 672 | // ----------------------------------------------------------------------------------- |
| 673 | |
| 674 | void compute_adam_update ( |
| 675 | size_t begin, |
| 676 | size_t end, |
| 677 | tensor& s, |
| 678 | tensor& m, |
| 679 | tensor& v, |
| 680 | const float t, |
| 681 | const float learning_rate, |
| 682 | const float weight_decay, |
| 683 | const float momentum1, |
| 684 | const float momentum2, |
| 685 | const tensor& params, |
| 686 | const tensor& params_grad |
| 687 | ) |
| 688 | { |
| 689 | DLIB_CASSERT(s.size() == m.size() && |
| 690 | s.size() == v.size() && |
| 691 | s.size() == params.size() && |
| 692 | s.size() == params_grad.size()); |
| 693 | DLIB_CASSERT(begin <= end && end <= params.size()); |
| 694 | const float eps = 1e-8; |
| 695 | const float alpha = learning_rate*std::sqrt(1-std::pow(momentum2,t))/(1-std::pow(momentum1, t)); |
| 696 | |
| 697 | // The loop is equivalent to doing this: |
| 698 | // m = momentum1*m + (1-momentum1) * (weight_decay*params + params_grad); |
| 699 | // v = momentum2*v + (1-momentum2)*squared(weight_decay*params + params_grad); |
| 700 | // s = -alpha*m/(sqrt(v) + eps); |
| 701 | auto pm = m.host(); |
| 702 | auto pv = v.host(); |
| 703 | auto ps = s.host_write_only(); |
| 704 | auto pparams = params.host(); |
| 705 | auto ppgrad = params_grad.host(); |
| 706 | for (size_t i = begin; i < end; ++i) |
| 707 | { |
| 708 | float g = weight_decay*pparams[i] + ppgrad[i]; |
| 709 | pm[i] = momentum1*pm[i] + (1-momentum1)*g; |
| 710 | pv[i] = momentum2*pv[i] + (1-momentum2)*g*g; |
| 711 | ps[i] = -alpha*pm[i]/(std::sqrt(pv[i]) + eps); |
| 712 | } |
| 713 | } |
| 714 | |
| 715 | // ----------------------------------------------------------------------------------- |
| 716 | |