sgd_op.h 5.1 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
Q
Qiao Longfei 已提交
2 3 4 5 6 7 8 9 10 11 12 13 14 15

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#pragma once
Y
Yi Wang 已提交
16 17 18
#include "paddle/fluid/framework/eigen.h"
#include "paddle/fluid/framework/op_registry.h"
#include "paddle/fluid/framework/selected_rows.h"
Q
Qiao Longfei 已提交
19 20 21 22

namespace paddle {
namespace operators {

C
chengduoZH 已提交
23
template <typename T>
Y
Yu Yang 已提交
24
class SGDOpKernel : public framework::OpKernel<T> {
25
 public:
26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67
  void Compute(const framework::ExecutionContext &ctx) const override {
    const auto *learning_rate = ctx.Input<framework::Tensor>("LearningRate");

    const auto *param_var = ctx.InputVar("Param");
    const auto *grad_var = ctx.InputVar("Grad");

    if (param_var->IsType<framework::LoDTensor>()) {
      const auto *param = ctx.Input<framework::Tensor>("Param");
      auto *param_out = ctx.Output<framework::Tensor>("ParamOut");

      // Actually, all tensors are LoDTensor except SelectedRows.
      if (grad_var->IsType<framework::LoDTensor>()) {
        param_out->mutable_data<T>(ctx.GetPlace());
        const auto *grad = ctx.Input<framework::Tensor>("Grad");

        auto p = framework::EigenVector<T>::Flatten(*param);
        auto g = framework::EigenVector<T>::Flatten(*grad);
        auto o = framework::EigenVector<T>::Flatten(*param_out);
        auto *lr = learning_rate->data<T>();

        o = p - lr[0] * g;
      } else if (grad_var->IsType<framework::SelectedRows>()) {
        // TODO(qijun): In Sparse SGD operator, in-place update is enforced.
        // This manual optimization brings difficulty to track data dependency.
        // It's better to find a more elegant solution.
        PADDLE_ENFORCE_EQ(param, param_out);
        const auto *grad = ctx.Input<framework::SelectedRows>("Grad");

        // for distributed training, a sparse var may be empty,
        // just skip updating.
        if (grad->rows().size() == 0) {
          return;
        }

        auto grad_height = grad->height();
        auto out_dims = param_out->dims();
        PADDLE_ENFORCE_EQ(grad_height, out_dims[0]);

        auto &grad_value = grad->value();
        auto &grad_rows = grad->rows();

        size_t grad_row_numel = grad_value.numel() / grad_rows.size();
A
Abhinav Arora 已提交
68 69
        PADDLE_ENFORCE_EQ(static_cast<int64_t>(grad_row_numel),
                          param_out->numel() / grad_height);
70 71 72 73 74 75 76

        auto *grad_data = grad_value.data<T>();
        auto *out_data = param_out->data<T>();
        auto *lr = learning_rate->data<T>();
        for (size_t i = 0; i < grad_rows.size(); i++) {
          PADDLE_ENFORCE(grad_rows[i] < grad_height,
                         "Input rows index should less than height");
A
Abhinav Arora 已提交
77
          for (size_t j = 0; j < grad_row_numel; j++) {
78 79 80 81 82 83 84 85 86 87 88 89 90 91
            out_data[grad_rows[i] * grad_row_numel + j] -=
                lr[0] * grad_data[i * grad_row_numel + j];
          }
        }
      } else {
        PADDLE_THROW("Unsupported Variable Type of Grad");
      }
    } else if (param_var->IsType<framework::SelectedRows>()) {
      PADDLE_ENFORCE(grad_var->IsType<framework::SelectedRows>(),
                     "when param "
                     "is SelectedRows, gradient should also be SelectedRows");
      const auto &param = param_var->Get<framework::SelectedRows>();
      auto *param_out = ctx.Output<framework::SelectedRows>("ParamOut");
      const auto &grad = grad_var->Get<framework::SelectedRows>();
C
chengduoZH 已提交
92

93 94
      // for distributed training, a sparse var may be empty,
      // just skip updating.
95
      if (grad.rows().size() == 0) {
96 97 98
        return;
      }

Q
qiaolongfei 已提交
99 100
      auto param_row_width = param.value().dims()[1];
      auto grad_row_width = grad.value().dims()[1];
101 102 103 104
      VLOG(40) << " param rows: " << param.rows().size()
               << " param memory rows: " << param.value().dims()[0]
               << " grad rows: " << grad.rows().size()
               << " grad memory rows: " << grad.value().dims()[0];
105 106
      PADDLE_ENFORCE_EQ(param_row_width, grad_row_width,
                        "param_row should have the same size with grad_row");
C
chengduoZH 已提交
107

108 109 110 111
      const auto *lr = learning_rate->data<T>();
      const auto *grad_data = grad.value().data<T>();
      auto *out_data = param_out->mutable_value()->data<T>();
      for (size_t i = 0; i < grad.rows().size(); i++) {
112
        int64_t id_index = param_out->AutoGrownIndex(grad.rows()[i], false);
Y
update  
Yancey1989 已提交
113 114
        PADDLE_ENFORCE_GE(id_index, static_cast<int64_t>(0),
                          "id should be in the table");
115
        for (int64_t j = 0; j < grad_row_width; j++) {
116 117
          out_data[id_index * grad_row_width + j] -=
              lr[0] * grad_data[i * grad_row_width + j];
C
chengduoZH 已提交
118 119
        }
      }
Q
qijun 已提交
120
    } else {
121
      PADDLE_THROW("Unsupported Variable Type of Parameter");
Q
qijun 已提交
122
    }
Q
Qiao Longfei 已提交
123 124 125 126
  }
};
}  // namespace operators
}  // namespace paddle