lookup_table_op.h 3.8 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserve.
   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

14
#include "paddle/framework/eigen.h"
15
#include "paddle/framework/lod_tensor.h"
16
#include "paddle/framework/op_registry.h"
17
#include "paddle/framework/selected_rows.h"
18 19 20 21

namespace paddle {
namespace operators {

F
fengjiayi 已提交
22
using LoDTensor = framework::LoDTensor;
23
using SelectedRows = framework::SelectedRows;
24 25

template <typename T>
Y
Yu Yang 已提交
26
class LookupTableKernel : public framework::OpKernel<T> {
27 28
 public:
  void Compute(const framework::ExecutionContext& context) const override {
F
fengjiayi 已提交
29 30 31
    auto* table_t = context.Input<LoDTensor>("W");      // float tensor
    auto* ids_t = context.Input<LoDTensor>("Ids");      // int tensor
    auto* output_t = context.Output<LoDTensor>("Out");  // float tensor
32

F
fengjiayi 已提交
33 34
    int N = table_t->dims()[0];
    int D = table_t->dims()[1];
F
fengjiayi 已提交
35 36 37
    auto* ids = ids_t->data<int64_t>();
    auto* table = table_t->data<T>();
    auto* output = output_t->mutable_data<T>(context.GetPlace());
38
    for (int64_t i = 0; i < ids_t->numel(); ++i) {
39 40 41 42 43 44 45 46
      PADDLE_ENFORCE_LT(ids[i], N);
      PADDLE_ENFORCE_GE(ids[i], 0);
      memcpy(output + i * D, table + ids[i] * D, D * sizeof(T));
    }
  }
};

template <typename T>
Y
Yu Yang 已提交
47
class LookupTableGradKernel : public framework::OpKernel<T> {
48 49
 public:
  void Compute(const framework::ExecutionContext& context) const override {
50 51
    bool is_sparse = context.Attr<bool>("is_sparse");
    if (is_sparse) {
F
fengjiayi 已提交
52 53 54
      auto* ids = context.Input<LoDTensor>("Ids");
      auto* table = context.Input<LoDTensor>("W");
      auto* d_output = context.Input<LoDTensor>(framework::GradVarName("Out"));
55
      auto* d_table = context.Output<SelectedRows>(framework::GradVarName("W"));
56

57 58
      auto* ids_data = ids->data<int64_t>();
      auto ids_dim = ids->dims();
59

60 61 62 63 64 65
      framework::Vector<int64_t> new_rows;
      new_rows.reserve(ids_dim[0]);
      for (int64_t i = 0; i < ids_dim[0]; i++) {
        new_rows.push_back(ids_data[i]);
      }
      d_table->set_rows(new_rows);
66

67 68 69 70 71 72 73 74 75 76 77 78
      auto* d_table_value = d_table->mutable_value();
      d_table_value->Resize({ids_dim[0], table->dims()[1]});
      d_table_value->mutable_data<T>(context.GetPlace());

      d_table->set_height(table->dims()[0]);

      auto* d_output_data = d_output->data<T>();
      auto* d_table_data = d_table_value->data<T>();

      PADDLE_ENFORCE_EQ(d_table_value->dims(), d_output->dims());
      memcpy(d_table_data, d_output_data, sizeof(T) * d_output->numel());
    } else {
F
fengjiayi 已提交
79 80 81 82
      auto* ids = context.Input<LoDTensor>("Ids");
      auto* d_output = context.Input<LoDTensor>(framework::GradVarName("Out"));
      auto* d_table = context.Output<LoDTensor>(framework::GradVarName("W"));
      auto* table = context.Input<LoDTensor>("W");
83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98

      auto* ids_data = ids->data<int64_t>();
      auto ids_dim = ids->dims();

      int N = table->dims()[0];
      int D = d_output->dims()[1];

      auto* d_output_data = d_output->data<T>();
      auto* d_table_data = d_table->mutable_data<T>(context.GetPlace());

      for (int64_t i = 0; i < ids->numel(); ++i) {
        PADDLE_ENFORCE_LT(ids_data[i], N);
        PADDLE_ENFORCE_GE(ids_data[i], 0);
        for (int j = 0; j < D; ++j) {
          d_table_data[ids_data[i] * D + j] = d_output_data[i * D + j];
        }
99 100 101 102 103 104 105
      }
    }
  }
};

}  // namespace operators
}  // namespace paddle