| #pragma once |
| |
| #include "caffe2/core/operator.h" |
| #include "caffe2/perfkernels/adagrad.h" |
| |
| namespace caffe2 { |
| |
| template <typename Context> |
| void adagrad_update( |
| int N, |
| const float* w, |
| const float* g, |
| const float* h, |
| float* nw, |
| float* nh, |
| float epsilon, |
| float decay, |
| const float* lr, |
| Context* /*context*/) { |
| return adagrad_update(N, w, g, h, nw, nh, epsilon, decay, lr[0]); |
| } |
| |
| template <typename Context> |
| void adagrad_update_output_effective_lr( |
| int N, |
| const float* paramIn, |
| const float* gradIn, |
| const float* momentIn, |
| float* paramOut, |
| float* momentOut, |
| float* effectiveLROut, |
| float epsilon, |
| float decay, |
| const float* lr, |
| Context* /*context*/) { |
| for (auto i = 0; i < N; ++i) { |
| float grad = gradIn[i]; |
| float moment = momentOut[i] = decay * momentIn[i] + grad * grad; |
| float effective_lr = effectiveLROut[i] = |
| lr[0] / (std::sqrt(moment) + epsilon); |
| paramOut[i] = paramIn[i] + effective_lr * grad; |
| } |
| } |
| |
| template <typename Context> |
| void adagrad_update_output_effective_lr_and_update( |
| int N, |
| const float* paramIn, |
| const float* gradIn, |
| const float* momentIn, |
| float* paramOut, |
| float* momentOut, |
| float* effectiveLROut, |
| float* updateOut, |
| float epsilon, |
| float decay, |
| const float* lr, |
| Context* /*context*/) { |
| for (auto i = 0; i < N; ++i) { |
| float grad = gradIn[i]; |
| float moment = momentOut[i] = decay * momentIn[i] + grad * grad; |
| float effective_lr = effectiveLROut[i] = |
| lr[0] / (std::sqrt(moment) + epsilon); |
| float update = updateOut[i] = effective_lr * grad; |
| paramOut[i] = paramIn[i] + update; |
| } |
| } |
| |
| template <typename T, class Context> |
| class AdagradOp final : public Operator<Context> { |
| public: |
| USE_OPERATOR_CONTEXT_FUNCTIONS; |
| AdagradOp(const OperatorDef& operator_def, Workspace* ws) |
| : Operator<Context>(operator_def, ws), |
| epsilon_(this->template GetSingleArgument<T>("epsilon", 1e-5f)), |
| decay_(this->template GetSingleArgument<T>("decay", 1.0f)) {} |
| |
| bool RunOnDevice() override { |
| CAFFE_ENFORCE_EQ( |
| Input(GRAD).numel(), |
| Input(MOMENT_1).numel(), |
| "PARAM size: ", |
| Input(PARAM).numel(), |
| ", GRAD size: ", |
| Input(GRAD).numel(), |
| ", MOMENT_1 size: ", |
| Input(MOMENT_1).numel(), |
| ", LR size: ", |
| Input(LR).numel()); |
| |
| CAFFE_ENFORCE_EQ(Input(GRAD).numel(), Input(PARAM).numel()); |
| Output(OUTPUT_PARAM)->ResizeLike(Input(PARAM)); |
| Output(OUTPUT_MOMENT_1)->ResizeLike(Input(MOMENT_1)); |
| if (OutputSize() == 2) { |
| adagrad_update<Context>( |
| Input(GRAD).numel(), |
| Input(PARAM).template data<T>(), |
| Input(GRAD).template data<T>(), |
| Input(MOMENT_1).template data<T>(), |
| Output(OUTPUT_PARAM)->template mutable_data<T>(), |
| Output(OUTPUT_MOMENT_1)->template mutable_data<T>(), |
| epsilon_, |
| decay_, |
| Input(LR).template data<T>(), |
| &context_); |
| } else if (OutputSize() == 3) { |
| Output(OUTPUT_EFFECTIVE_LR)->ResizeLike(Input(GRAD)); |
| adagrad_update_output_effective_lr<Context>( |
| Input(GRAD).numel(), |
| Input(PARAM).template data<T>(), |
| Input(GRAD).template data<T>(), |
| Input(MOMENT_1).template data<T>(), |
| Output(OUTPUT_PARAM)->template mutable_data<T>(), |
| Output(OUTPUT_MOMENT_1)->template mutable_data<T>(), |
| Output(OUTPUT_EFFECTIVE_LR)->template mutable_data<T>(), |
| epsilon_, |
| decay_, |
| Input(LR).template data<T>(), |
| &context_); |
| } else { |
| Output(OUTPUT_EFFECTIVE_LR)->ResizeLike(Input(GRAD)); |
| Output(OUTPUT_UPDATE)->ResizeLike(Input(GRAD)); |
| adagrad_update_output_effective_lr_and_update<Context>( |
| Input(GRAD).numel(), |
| Input(PARAM).template data<T>(), |
| Input(GRAD).template data<T>(), |
| Input(MOMENT_1).template data<T>(), |
| Output(OUTPUT_PARAM)->template mutable_data<T>(), |
| Output(OUTPUT_MOMENT_1)->template mutable_data<T>(), |
| Output(OUTPUT_EFFECTIVE_LR)->template mutable_data<T>(), |
| Output(OUTPUT_UPDATE)->template mutable_data<T>(), |
| epsilon_, |
| decay_, |
| Input(LR).template data<T>(), |
| &context_); |
| } |
| |
| return true; |
| } |
| |
| protected: |
| T epsilon_; |
| T decay_; |
| INPUT_TAGS(PARAM, MOMENT_1, GRAD, LR); |
| OUTPUT_TAGS( |
| OUTPUT_PARAM, |
| OUTPUT_MOMENT_1, |
| OUTPUT_EFFECTIVE_LR, |
| OUTPUT_UPDATE); |
| }; |
| |
| template <typename T, class Context> |
| class SparseAdagradOp final : public Operator<Context> { |
| public: |
| USE_OPERATOR_CONTEXT_FUNCTIONS; |
| SparseAdagradOp(const OperatorDef& operator_def, Workspace* ws) |
| : Operator<Context>(operator_def, ws), |
| epsilon_(this->template GetSingleArgument<float>("epsilon", 1e-5f)) {} |
| |
| bool RunOnDevice() override { |
| // Enforce shapes |
| CAFFE_ENFORCE_EQ(Input(PARAM).numel(), Input(MOMENT_1).numel()); |
| CAFFE_ENFORCE_EQ(Input(LR).numel(), 1); |
| CAFFE_ENFORCE_EQ( |
| Input(PARAM).size_from_dim(1), |
| Input(GRAD).size_from_dim(Input(INDICES).dim())); |
| |
| return DispatchHelper<TensorTypes<int32_t, int64_t>>::call( |
| this, Input(INDICES)); |
| } |
| |
| template <typename SIndex> |
| bool DoRunWithType() { |
| const auto* lr = Input(LR).template data<T>(); |
| const auto* indices = Input(INDICES).template data<SIndex>(); |
| const auto* gradIn = Input(GRAD).template data<T>(); |
| const auto* paramIn = Input(PARAM).template data<T>(); |
| const auto* momentIn = Input(MOMENT_1).template data<T>(); |
| auto* paramOut = Output(OUTPUT_PARAM)->template mutable_data<T>(); |
| auto* momentOut = Output(OUTPUT_MOMENT_1)->template mutable_data<T>(); |
| |
| auto n = Input(INDICES).numel(); |
| if (n == 0) { |
| return true; |
| } |
| |
| auto block_size = Input(GRAD).numel() / n; |
| for (auto i = 0; i < n; ++i) { |
| auto idx = indices[i]; |
| if (block_size == 1) { |
| float gi = gradIn[i]; |
| float hi = momentOut[idx] = momentIn[idx] + gi * gi; |
| paramOut[idx] = paramIn[idx] + lr[0] * gi / (std::sqrt(hi) + epsilon_); |
| } else { |
| auto offsetI = i * block_size; |
| auto offsetIdx = idx * block_size; |
| |
| #ifndef NDEBUG |
| CAFFE_ENFORCE_GE( |
| Input(PARAM).numel(), |
| block_size + offsetIdx, |
| this->debug_def().input(PARAM), |
| ", out of bound, idx:", |
| idx, |
| " for input i:", |
| i, |
| " and block size:", |
| block_size); |
| CAFFE_ENFORCE_GE( |
| Input(GRAD).numel(), |
| block_size + offsetI, |
| this->debug_def().input(GRAD), |
| ", out of bound idx, idx:", |
| idx, |
| " for input i:", |
| i); |
| #endif |
| adagrad_update( |
| block_size, |
| paramIn + offsetIdx, |
| gradIn + offsetI, |
| momentIn + offsetIdx, |
| paramOut + offsetIdx, |
| momentOut + offsetIdx, |
| epsilon_, |
| 1.0f, |
| lr, |
| &context_); |
| } |
| } |
| return true; |
| } |
| |
| protected: |
| T epsilon_; |
| INPUT_TAGS(PARAM, MOMENT_1, INDICES, GRAD, LR); |
| OUTPUT_TAGS(OUTPUT_PARAM, OUTPUT_MOMENT_1); |
| }; |
| |
| template <typename T, class Context> |
| class RowWiseSparseAdagradOp final : public Operator<Context> { |
| public: |
| USE_OPERATOR_CONTEXT_FUNCTIONS; |
| RowWiseSparseAdagradOp(const OperatorDef& operator_def, Workspace* ws) |
| : Operator<Context>(operator_def, ws), |
| epsilon_(this->template GetSingleArgument<float>("epsilon", 1e-5f)) {} |
| |
| bool RunOnDevice() override { |
| // Enforce shapes |
| CAFFE_ENFORCE_EQ(Input(PARAM).sizes()[0], Input(MOMENT_1).numel()); |
| CAFFE_ENFORCE_EQ(Input(LR).numel(), 1); |
| CAFFE_ENFORCE_EQ( |
| Input(PARAM).size_from_dim(1), |
| Input(GRAD).size_from_dim(Input(INDICES).dim())); |
| |
| return DispatchHelper<TensorTypes<int32_t, int64_t>>::call( |
| this, Input(INDICES)); |
| } |
| |
| template <typename SIndex> |
| bool DoRunWithType() { |
| const auto* lr = Input(LR).template data<T>(); |
| const auto* indices = Input(INDICES).template data<SIndex>(); |
| const auto* gradIn = Input(GRAD).template data<T>(); |
| const auto* paramIn = Input(PARAM).template data<T>(); |
| const auto* momentIn = Input(MOMENT_1).template data<T>(); |
| auto* paramOut = Output(OUTPUT_PARAM)->template mutable_data<T>(); |
| auto* momentOut = Output(OUTPUT_MOMENT_1)->template mutable_data<T>(); |
| |
| auto n = Input(INDICES).numel(); |
| if (n == 0) { |
| return true; |
| } |
| |
| auto block_size = Input(GRAD).numel() / n; |
| |
| for (auto i = 0; i < n; ++i) { |
| auto idx = indices[i]; |
| if (block_size == 1) { |
| float gi = gradIn[i]; |
| float hi = momentOut[idx] = momentIn[idx] + gi * gi; |
| paramOut[idx] = paramIn[idx] + lr[0] * gi / (std::sqrt(hi) + epsilon_); |
| } else { |
| auto offsetI = i * block_size; |
| auto offsetIdx = idx * block_size; |
| |
| #ifndef NDEBUG |
| CAFFE_ENFORCE_GE( |
| Input(PARAM).numel(), |
| block_size + offsetIdx, |
| this->debug_def().input(PARAM), |
| ", out of bound, idx:", |
| idx, |
| " for input i:", |
| i, |
| " and block size:", |
| block_size); |
| CAFFE_ENFORCE_GE( |
| Input(GRAD).numel(), |
| block_size + offsetI, |
| this->debug_def().input(GRAD), |
| ", out of bound idx, idx:", |
| idx, |
| " for input i:", |
| i); |
| #endif |
| |
| const float* w = paramIn + offsetIdx; |
| const float* g = gradIn + offsetI; |
| const float* h = momentIn + idx; |
| float* nw = paramOut + offsetIdx; |
| float* nh = momentOut + idx; |
| float hs = 0.; |
| for (auto j = 0; j < block_size; ++j) { |
| float gj = g[j]; |
| hs += gj * gj; |
| } |
| float hi = nh[0] = h[0] + hs / block_size; |
| float step = lr[0] / (std::sqrt(hi) + epsilon_); |
| for (auto j = 0; j < block_size; ++j) { |
| nw[j] = w[j] + g[j] * step; |
| } |
| } |
| } |
| return true; |
| } |
| |
| protected: |
| T epsilon_; |
| INPUT_TAGS(PARAM, MOMENT_1, INDICES, GRAD, LR); |
| OUTPUT_TAGS(OUTPUT_PARAM, OUTPUT_MOMENT_1); |
| }; |
| } |
