Refine lookup_table_op (#5257)

1. Change some `auto` to `auto*` 2. Change `Tensor` to `LoDTensor`

Author: JiayiFeng

paddle/operators/lookup_table_op.cc

@@ -43,7 +43,7 @@ class LookupTableOp : public framework::OperatorWithKernel {
  protected:
  framework::DataType IndicateDataType(
  const framework::ExecutionContext& ctx) const override {
- return framework::ToDataType(ctx.Input<Tensor>("W")->type());
+ return framework::ToDataType(ctx.Input<LoDTensor>("W")->type());
  }
 };
 
@@ -93,7 +93,7 @@ class LookupTableOpGrad : public framework::OperatorWithKernel {
  protected:
  framework::DataType IndicateDataType(
  const framework::ExecutionContext& ctx) const override {
- return framework::ToDataType(ctx.Input<Tensor>("W")->type());
+ return framework::ToDataType(ctx.Input<LoDTensor>("W")->type());
  }
 };
 

paddle/operators/lookup_table_op.cu

@@ -61,16 +61,16 @@ template <typename T>
 class LookupTableCUDAKernel : public framework::OpKernel<T> {
  public:
  void Compute(const framework::ExecutionContext& context) const override {
- auto table_t = context.Input<Tensor>("W");
- auto ids_t = context.Input<Tensor>("Ids");
- auto output_t = context.Output<Tensor>("Out");
+ auto* table_t = context.Input<LoDTensor>("W");
+ auto* ids_t = context.Input<LoDTensor>("Ids");
+ auto* output_t = context.Output<LoDTensor>("Out");
 
  size_t N = table_t->dims()[0];
  size_t D = table_t->dims()[1];
  size_t K = ids_t->numel();
- auto ids = ids_t->data<int64_t>();
- auto table = table_t->data<T>();
- auto output = output_t->mutable_data<T>(context.GetPlace());
+ auto* ids = ids_t->data<int64_t>();
+ auto* table = table_t->data<T>();
+ auto* output = output_t->mutable_data<T>(context.GetPlace());
 
  dim3 threads(128, 8);
  dim3 grids(8, 1);
@@ -87,9 +87,9 @@ class LookupTableGradCUDAKernel : public framework::OpKernel<T> {
  void Compute(const framework::ExecutionContext& context) const override {
  bool is_sparse = context.Attr<bool>("is_sparse");
  if (is_sparse) {
- auto* ids = context.Input<Tensor>("Ids");
- auto* table = context.Input<Tensor>("W");
- auto* d_output = context.Input<Tensor>(framework::GradVarName("Out"));
+ auto* ids = context.Input<LoDTensor>("Ids");
+ auto* table = context.Input<LoDTensor>("W");
+ auto* d_output = context.Input<LoDTensor>(framework::GradVarName("Out"));
  auto* d_table = context.Output<SelectedRows>(framework::GradVarName("W"));
 
  auto* ids_data = ids->data<int64_t>();
@@ -119,9 +119,9 @@ class LookupTableGradCUDAKernel : public framework::OpKernel<T> {
  d_output->numel(), stream);
 
  } else {
- auto ids_t = context.Input<Tensor>("Ids");
- auto d_output_t = context.Input<Tensor>(framework::GradVarName("Out"));
- auto d_table_t = context.Output<Tensor>(framework::GradVarName("W"));
+ auto ids_t = context.Input<LoDTensor>("Ids");
+ auto d_output_t = context.Input<LoDTensor>(framework::GradVarName("Out"));
+ auto d_table_t = context.Output<LoDTensor>(framework::GradVarName("W"));
 
  int N = d_table_t->dims()[0];
  int D = d_table_t->dims()[1];

paddle/operators/lookup_table_op.h

@@ -19,22 +19,22 @@
 namespace paddle {
 namespace operators {
 
-using Tensor = framework::Tensor;
+using LoDTensor = framework::LoDTensor;
 using SelectedRows = framework::SelectedRows;
 
 template <typename T>
 class LookupTableKernel : public framework::OpKernel<T> {
  public:
  void Compute(const framework::ExecutionContext& context) const override {
- auto table_t = context.Input<Tensor>("W"); // float tensor
- auto ids_t = context.Input<Tensor>("Ids"); // int tensor
- auto output_t = context.Output<Tensor>("Out"); // float tensor
+ 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
 
  int N = table_t->dims()[0];
  int D = table_t->dims()[1];
- auto ids = ids_t->data<int64_t>();
- auto table = table_t->data<T>();
- auto output = output_t->mutable_data<T>(context.GetPlace());
+ auto* ids = ids_t->data<int64_t>();
+ auto* table = table_t->data<T>();
+ auto* output = output_t->mutable_data<T>(context.GetPlace());
  for (int64_t i = 0; i < ids_t->numel(); ++i) {
  PADDLE_ENFORCE_LT(ids[i], N);
  PADDLE_ENFORCE_GE(ids[i], 0);
@@ -49,9 +49,9 @@ class LookupTableGradKernel : public framework::OpKernel<T> {
  void Compute(const framework::ExecutionContext& context) const override {
  bool is_sparse = context.Attr<bool>("is_sparse");
  if (is_sparse) {
- auto* ids = context.Input<Tensor>("Ids");
- auto* table = context.Input<Tensor>("W");
- auto* d_output = context.Input<Tensor>(framework::GradVarName("Out"));
+ auto* ids = context.Input<LoDTensor>("Ids");
+ auto* table = context.Input<LoDTensor>("W");
+ auto* d_output = context.Input<LoDTensor>(framework::GradVarName("Out"));
  auto* d_table = context.Output<SelectedRows>(framework::GradVarName("W"));
 
  auto* ids_data = ids->data<int64_t>();
@@ -76,10 +76,10 @@ class LookupTableGradKernel : public framework::OpKernel<T> {
  PADDLE_ENFORCE_EQ(d_table_value->dims(), d_output->dims());
  memcpy(d_table_data, d_output_data, sizeof(T) * d_output->numel());
  } else {
- auto* ids = context.Input<Tensor>("Ids");
- auto* d_output = context.Input<Tensor>(framework::GradVarName("Out"));
- auto* d_table = context.Output<Tensor>(framework::GradVarName("W"));
- auto* table = context.Input<Tensor>("W");
+ 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");
 
  auto* ids_data = ids->data<int64_t>();
  auto ids_dim = ids->dims();