From 92cc33c0e48a98c07a6cf686a22f15bf396fe701 Mon Sep 17 00:00:00 2001 From: Chengmo Date: Thu, 9 Apr 2020 14:00:07 +0800 Subject: [PATCH] Cherry-pick index sample op in contrib (#23522) * test=develop, add index sample op in contrib --- paddle/fluid/operators/index_sample_op.cc | 158 +++++++++++++++ paddle/fluid/operators/index_sample_op.h | 183 ++++++++++++++++++ python/paddle/fluid/contrib/layers/nn.py | 63 +++++- .../tests/unittests/test_index_sample_op.py | 127 ++++++++++++ 4 files changed, 530 insertions(+), 1 deletion(-) create mode 100644 paddle/fluid/operators/index_sample_op.cc create mode 100644 paddle/fluid/operators/index_sample_op.h create mode 100644 python/paddle/fluid/tests/unittests/test_index_sample_op.py diff --git a/paddle/fluid/operators/index_sample_op.cc b/paddle/fluid/operators/index_sample_op.cc new file mode 100644 index 00000000000..ad13b0d95e9 --- /dev/null +++ b/paddle/fluid/operators/index_sample_op.cc @@ -0,0 +1,158 @@ +/* Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. + +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. */ + +#include "paddle/fluid/operators/index_sample_op.h" +#include +#include +#include "paddle/fluid/framework/no_need_buffer_vars_inference.h" +#include "paddle/fluid/framework/op_registry.h" +#include "paddle/fluid/platform/enforce.h" + +namespace paddle { +namespace operators { +class IndexSampleOpMaker : public framework::OpProtoAndCheckerMaker { + public: + void Make() override { + AddInput("X", "Input(Tensor), dtype support int32/int64/float/double"); + AddInput("Index", "Index(Tensor), dtype support int32/int64"); + AddOutput("Out", "Return the element of input at index"); + + AddComment(R"DOC( + IndexSample OP returns the element of the specified location of X, + and the location is specified by Index. + + X tensor and Index tensor's shape must be 2-D, + dimension at 0 which usually is batch size must be equal. + + The returned tensor has the same shape and dimensions as the Index tensor. + )DOC"); + } +}; + +class IndexSampleOp : public framework::OperatorWithKernel { + public: + using framework::OperatorWithKernel::OperatorWithKernel; + void InferShape(framework::InferShapeContext* ctx) const override { + PADDLE_ENFORCE_EQ(ctx->HasInput("X"), true, + platform::errors::InvalidArgument( + "Inputs(Input) of FindByIndex should not be null.")); + PADDLE_ENFORCE_EQ(ctx->HasInput("Index"), true, + platform::errors::InvalidArgument( + "Inputs(Index) of FindByIndex should not be null.")); + + auto input_dims = ctx->GetInputDim("X"); + PADDLE_ENFORCE_EQ( + input_dims.size(), 2, + platform::errors::InvalidArgument( + "Inputs(X) shape of IndexSample op should be 2-D, but " + "got X's shape = [%s], please check X shape.", + input_dims)); + + auto index_dims = ctx->GetInputDim("Index"); + PADDLE_ENFORCE_EQ( + input_dims.size(), 2, + platform::errors::InvalidArgument( + "Inputs(Index) shape of IndexSample op should be 2-D, but " + "got Index's shape [%s] , please check index shape.", + input_dims)); + if (ctx->IsRuntime()) { + PADDLE_ENFORCE_EQ(input_dims[0], index_dims[0], + platform::errors::InvalidArgument( + "Inputs(X)'s value of dimension 0 must same with " + "Inputs(Index)'s value of dimension 0, but " + "got %d of Inputs(X), and got %d of Inputs(Index), " + "please check Inputs shape.", + input_dims[0], index_dims[0])); + } + ctx->SetOutputDim("Out", index_dims); + auto type = ctx->GetInputsVarType("Index")[0]; + if (type == framework::proto::VarType::LOD_TENSOR) { + ctx->ShareLoD("Index", /*->*/ "Out"); + } + } + + protected: + framework::OpKernelType GetExpectedKernelType( + const framework::ExecutionContext& ctx) const override { + auto data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X"); + return framework::OpKernelType(data_type, ctx.device_context()); + } +}; + +class IndexSampleGradOp : public framework::OperatorWithKernel { + public: + using framework::OperatorWithKernel::OperatorWithKernel; + + void InferShape(framework::InferShapeContext* ctx) const override { + PADDLE_ENFORCE_EQ( + ctx->HasInput("Index"), true, + platform::errors::InvalidArgument("Input(Index) should be not null.")); + PADDLE_ENFORCE_EQ(ctx->HasInput(framework::GradVarName("Out")), true, + platform::errors::InvalidArgument( + "Input(Out@GRAD) should be not null.")); + PADDLE_ENFORCE_EQ(ctx->HasOutput(framework::GradVarName("X")), true, + platform::errors::InvalidArgument( + "Output(X@GRAD) should be not null.")); + + ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X")); + } + + protected: + framework::OpKernelType GetExpectedKernelType( + const framework::ExecutionContext& ctx) const override { + auto data_type = OperatorWithKernel::IndicateVarDataType( + ctx, framework::GradVarName("Out")); + return framework::OpKernelType(data_type, ctx.device_context()); + } +}; + +template +class IndexSampleGradMaker : public framework::SingleGradOpMaker { + public: + using framework::SingleGradOpMaker::SingleGradOpMaker; + + protected: + std::unique_ptr Apply() const override { + std::unique_ptr op(new T()); + op->SetType("index_sample_grad"); + op->SetInput("X", this->Input("X")); + op->SetInput("Index", this->Input("Index")); + op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out")); + op->SetOutput(framework::GradVarName("X"), this->InputGrad("X")); + return op; + } +}; + +DECLARE_NO_NEED_BUFFER_VARS_INFERENCE(IndexSampleGradNoNeedBufferVarInferer, + "X"); +} // namespace operators +} // namespace paddle + +namespace ops = paddle::operators; +REGISTER_OPERATOR(index_sample, ops::IndexSampleOp, ops::IndexSampleOpMaker, + ops::IndexSampleGradMaker, + ops::IndexSampleGradMaker); +REGISTER_OPERATOR(index_sample_grad, ops::IndexSampleGradOp, + ops::IndexSampleGradNoNeedBufferVarInferer); +REGISTER_OP_CPU_KERNEL( + index_sample, ops::IndexSampleKernel, + ops::IndexSampleKernel, + ops::IndexSampleKernel, + ops::IndexSampleKernel); +REGISTER_OP_CPU_KERNEL( + index_sample_grad, + ops::IndexSampleGradKernel, + ops::IndexSampleGradKernel, + ops::IndexSampleGradKernel, + ops::IndexSampleGradKernel); diff --git a/paddle/fluid/operators/index_sample_op.h b/paddle/fluid/operators/index_sample_op.h new file mode 100644 index 00000000000..a4895550a32 --- /dev/null +++ b/paddle/fluid/operators/index_sample_op.h @@ -0,0 +1,183 @@ +/* Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. +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 + +#include +#include +#include +#include +#include +#include +#include +#include "paddle/fluid/framework/op_registry.h" + +namespace paddle { +namespace operators { + +using Tensor = framework::Tensor; +using LoDTensor = framework::LoDTensor; +using DDim = framework::DDim; + +template +void IndexSampleInner(const framework::ExecutionContext &context, + const LoDTensor &input, const LoDTensor &index, + LoDTensor *output) { + auto input_dims = input.dims(); + auto index_dims = index.dims(); + + int batch_size = input_dims[0]; + auto value_length = input_dims[1]; + auto index_length = index_dims[1]; + int index_ids_num = index.numel(); + auto *input_data = input.data(); + auto *index_data = index.data(); + + std::vector res{}; + for (int i = 0; i < index_ids_num; i++) { + int b = floor(i / index_length); + PADDLE_ENFORCE_GE( + index_data[i], 0, + platform::errors::InvalidArgument( + "Variable value (index) of OP(index_sample) " + "expected >= 0 and < %ld, but got %ld. Please check input " + "value.", + value_length, index_data[i])); + PADDLE_ENFORCE_LT( + index_data[i], value_length, + platform::errors::InvalidArgument( + "Variable value (index) of OP(index_sample) " + "expected >= 0 and < %ld, but got %ld. Please check input " + "value.", + value_length, index_data[i])); + + int v_i = b * value_length + static_cast(index_data[i]); + T v = input_data[v_i]; + VLOG(4) << "Index Sample: batch = " << b << " index = " << v_i + << " value = " << v; + res.push_back(v); + } + + auto ddim = framework::make_ddim({batch_size, index_length}); + output->Resize(ddim); + T *out_data = output->mutable_data(context.GetPlace()); + + memcpy(out_data, &res[0], sizeof(T) * index_ids_num); +} + +template +class IndexSampleKernel : public framework::OpKernel { + public: + void Compute(const framework::ExecutionContext &ctx) const override { + auto *input_var = ctx.InputVar("X"); + auto *index_var = ctx.InputVar("Index"); + + auto &input_tensor = input_var->Get(); + auto &index_tensor = index_var->Get(); + + auto *out_var = ctx.OutputVar("Out"); + auto *out_tensor = out_var->GetMutable(); + + const auto &index_type = index_tensor.type(); + bool index_type_match = index_type == framework::proto::VarType::INT32 || + index_type == framework::proto::VarType::INT64; + PADDLE_ENFORCE_EQ(index_type_match, true, + platform::errors::InvalidArgument( + "Input(Index) holds the wrong type, it holds %s, but " + "desires to be %s or %s", + paddle::framework::DataTypeToString(index_type), + paddle::framework::DataTypeToString( + framework::proto::VarType::INT32), + paddle::framework::DataTypeToString( + framework::proto::VarType::INT64))); + if (index_type == framework::proto::VarType::INT32) { + IndexSampleInner(ctx, input_tensor, index_tensor, out_tensor); + } else if (index_type == framework::proto::VarType::INT64) { + IndexSampleInner(ctx, input_tensor, index_tensor, out_tensor); + } + } +}; + +template +void IndexSampleGradInner(const framework::ExecutionContext &context, + const LoDTensor &out_grad, const LoDTensor &index, + LoDTensor *x_grad) { + auto index_dims = index.dims(); + auto x_grad_dims = x_grad->dims(); + + int batch_size = x_grad_dims[0]; + auto value_length = x_grad_dims[1]; + auto index_length = index_dims[1]; + int index_ids_num = index.numel(); + + T *x_grad_data = x_grad->mutable_data(context.GetPlace()); + auto *out_grad_data = out_grad.data(); + auto *index_data = index.data(); + + memset(x_grad_data, 0, batch_size * value_length * sizeof(T)); + + for (int i = 0; i < index_ids_num; i++) { + int b = floor(i / index_length); + PADDLE_ENFORCE_GE( + index_data[i], 0, + platform::errors::InvalidArgument( + "Variable value (index) of OP(index_sample_grad) " + "expected >= 0 and < %ld, but got %ld. Please check input " + "value.", + value_length, index_data[i])); + PADDLE_ENFORCE_LT( + index_data[i], value_length, + platform::errors::InvalidArgument( + "Variable value (index) of OP(index_sample_grad) " + "expected >= 0 and < %ld, but got %ld. Please check input " + "value.", + value_length, index_data[i])); + int v_i = b * value_length + static_cast(index_data[i]); + x_grad_data[v_i] += out_grad_data[i]; + } +} + +template +class IndexSampleGradKernel : public framework::OpKernel { + public: + void Compute(const framework::ExecutionContext &context) const override { + auto *index_var = context.InputVar("Index"); + auto *x_grad_var = context.OutputVar(framework::GradVarName("X")); + auto *out_grad_var = context.InputVar(framework::GradVarName("Out")); + + auto &index_tensor = index_var->Get(); + auto &out_grad_tensor = out_grad_var->Get(); + auto *x_grad_tensor = x_grad_var->GetMutable(); + + const auto &index_type = index_tensor.type(); + bool index_type_match = index_type == framework::proto::VarType::INT32 || + index_type == framework::proto::VarType::INT64; + PADDLE_ENFORCE_EQ(index_type_match, true, + platform::errors::InvalidArgument( + "Input(Index) holds the wrong type, it holds %s, but " + "desires to be %s or %s", + paddle::framework::DataTypeToString(index_type), + paddle::framework::DataTypeToString( + framework::proto::VarType::INT32), + paddle::framework::DataTypeToString( + framework::proto::VarType::INT64))); + if (index_type == framework::proto::VarType::INT32) { + IndexSampleGradInner(context, out_grad_tensor, index_tensor, + x_grad_tensor); + } else if (index_type == framework::proto::VarType::INT64) { + IndexSampleGradInner(context, out_grad_tensor, index_tensor, + x_grad_tensor); + } + } +}; + +} // namespace operators +} // namespace paddle diff --git a/python/paddle/fluid/contrib/layers/nn.py b/python/paddle/fluid/contrib/layers/nn.py index 066e48db20e..61e2a96194c 100644 --- a/python/paddle/fluid/contrib/layers/nn.py +++ b/python/paddle/fluid/contrib/layers/nn.py @@ -25,7 +25,7 @@ from paddle.fluid.layer_helper import LayerHelper from paddle.fluid.layers import utils from ... import unique_name from paddle.fluid.initializer import Normal, Constant, NumpyArrayInitializer -from paddle.fluid.data_feeder import check_type, check_dtype, convert_dtype +from paddle.fluid.data_feeder import check_type_and_dtype, check_type, check_dtype, convert_dtype from paddle.fluid.framework import Variable, convert_np_dtype_to_dtype_ from paddle.fluid.layers import slice, reshape @@ -39,6 +39,7 @@ __all__ = [ 'multiclass_nms2', 'search_pyramid_hash', 'shuffle_batch', + 'index_sample', 'tdm_child', 'tdm_sampler', ] @@ -816,6 +817,66 @@ def shuffle_batch(x, seed=None): return out +def index_sample(x, index): + """ + **IndexSample Layer** + IndexSample OP returns the element of the specified location of X, + and the location is specified by Index. + + .. code-block:: text + + Given: + X = [[1, 2, 3, 4, 5], + [6, 7, 8, 9, 10]] + Index = [[0, 1, 3], + [0, 2, 4]] + Then: + Out = [[1, 2, 4], + [6, 8, 10]] + + Args: + x (Variable): The source input tensor with 2-D shape. Supported data type is + int32, int64, float32, float64. + index (Variable): The index input tensor with 2-D shape, first dimension should be same with X. + Data type is int32 or int64. + + Returns: + output (Variable): The output is a tensor with the same shape as index. + + Examples: + .. code-block:: python + import paddle.fluid as fluid + import numpy as np + # create x value + x_shape = (2, 5) + x_type = "float64" + x_np = np.random.random(x_shape).astype(x_type) + # create index value + index_shape = (2, 3) + index_type = "int32" + index_np = np.random.randint(low=0, + high=x_shape[1], + size=index_shape).astype(index_type) + x = fluid.data(name='x', shape=[-1, 5], dtype='float64') + index = fluid.data(name='index', shape=[-1, 3], dtype='int32') + output = fluid.contrib.layers.index_sample(x=x, index=index) + """ + helper = LayerHelper("index_sample", **locals()) + check_type_and_dtype(x, 'x', Variable, + ['float32', 'float64', 'int32', 'int64'], + 'fluid.contrib.layers.index_sample') + check_type_and_dtype(index, 'index', Variable, ['int32', 'int64'], + 'fluid.contrib.layers.index_sample') + out = helper.create_variable_for_type_inference(dtype=x.dtype) + + helper.append_op( + type='index_sample', + inputs={'X': x, + 'Index': index}, + outputs={'Out': out}) + return out + + def tdm_child(x, node_nums, child_nums, param_attr=None, dtype='int32'): """ **Tdm Child** diff --git a/python/paddle/fluid/tests/unittests/test_index_sample_op.py b/python/paddle/fluid/tests/unittests/test_index_sample_op.py new file mode 100644 index 00000000000..a898364e258 --- /dev/null +++ b/python/paddle/fluid/tests/unittests/test_index_sample_op.py @@ -0,0 +1,127 @@ +# Copyright (c) 2020 PaddlePaddle Authors. All Rights Reserved. +# +# 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. + +from __future__ import print_function + +import unittest +import numpy as np +from op_test import OpTest + + +class TestIndexSampleOp(OpTest): + def setUp(self): + self.op_type = "index_sample" + self.config() + xnp = np.random.random(self.x_shape).astype(self.x_type) + indexnp = np.random.randint( + low=0, high=self.x_shape[1], + size=self.index_shape).astype(self.index_type) + self.inputs = {'X': xnp, 'Index': indexnp} + index_array = [] + for i in range(self.index_shape[0]): + for j in indexnp[i]: + index_array.append(xnp[i, j]) + out = np.reshape(index_array, self.index_shape) + self.outputs = {'Out': out} + + def test_check_output(self): + self.check_output() + + def test_check_grad(self): + self.check_grad(['X'], 'Out') + + def config(self): + """ + For multi-dimension input + """ + self.x_shape = (10, 20) + self.x_type = "float64" + self.index_shape = (10, 10) + self.index_type = "int32" + + +class TestCase1(TestIndexSampleOp): + def config(self): + """ + For one dimension input + """ + self.x_shape = (100, 1) + self.x_type = "float64" + self.index_shape = (100, 1) + self.index_type = "int32" + + +class TestCase2(TestIndexSampleOp): + def config(self): + """ + For int64_t index type + """ + self.x_shape = (10, 100) + self.x_type = "float64" + self.index_shape = (10, 10) + self.index_type = "int64" + + +class TestCase3(TestIndexSampleOp): + def config(self): + """ + For int index type + """ + self.x_shape = (10, 100) + self.x_type = "float64" + self.index_shape = (10, 10) + self.index_type = "int32" + + +class TestCase4(TestIndexSampleOp): + def config(self): + """ + For int64 index type + """ + self.x_shape = (10, 100) + self.x_type = "float64" + self.index_shape = (10, 10) + self.index_type = "int64" + + +class TestIndexSampleShape(unittest.TestCase): + def test_shape(self): + import paddle.fluid as fluid + import paddle + + # create x value + x_shape = (2, 5) + x_type = "float64" + x_np = np.random.random(x_shape).astype(x_type) + + # create index value + index_shape = (2, 3) + index_type = "int32" + index_np = np.random.randint( + low=0, high=x_shape[1], size=index_shape).astype(index_type) + + x = fluid.data(name='x', shape=[-1, 5], dtype='float64') + index = fluid.data(name='index', shape=[-1, 3], dtype='int32') + output = fluid.contrib.layers.index_sample(x=x, index=index) + + place = fluid.CPUPlace() + exe = fluid.Executor(place=place) + exe.run(fluid.default_startup_program()) + + feed = {'x': x_np, 'index': index_np} + res = exe.run(feed=feed, fetch_list=[output]) + + +if __name__ == "__main__": + unittest.main() -- GitLab