From 1bfff02047c6b9db9e864864aba827fdcdc09bae Mon Sep 17 00:00:00 2001 From: zhoukunsheng Date: Fri, 10 May 2019 16:58:45 +0800 Subject: [PATCH] Add Diag Op(#17027) --- paddle/fluid/API.spec | 1 + paddle/fluid/operators/diag_op.cc | 60 +++++++++++++++++++ paddle/fluid/operators/diag_op.cu | 23 +++++++ paddle/fluid/operators/diag_op.h | 59 ++++++++++++++++++ python/paddle/fluid/layers/tensor.py | 38 +++++++++++- .../paddle/fluid/tests/unittests/test_diag.py | 43 +++++++++++++ 6 files changed, 223 insertions(+), 1 deletion(-) create mode 100644 paddle/fluid/operators/diag_op.cc create mode 100644 paddle/fluid/operators/diag_op.cu create mode 100644 paddle/fluid/operators/diag_op.h create mode 100644 python/paddle/fluid/tests/unittests/test_diag.py diff --git a/paddle/fluid/API.spec b/paddle/fluid/API.spec index 354be97f28..f285a5d60b 100644 --- a/paddle/fluid/API.spec +++ b/paddle/fluid/API.spec @@ -270,6 +270,7 @@ paddle.fluid.layers.isfinite (ArgSpec(args=['x'], varargs=None, keywords=None, d paddle.fluid.layers.range (ArgSpec(args=['start', 'end', 'step', 'dtype'], varargs=None, keywords=None, defaults=None), ('document', '2ec937ede953ded2fdff2675883900bb')) paddle.fluid.layers.linspace (ArgSpec(args=['start', 'stop', 'num', 'dtype'], varargs=None, keywords=None, defaults=None), ('document', '495e21e9a848c2d075a102802fc67756')) paddle.fluid.layers.zeros_like (ArgSpec(args=['x', 'out'], varargs=None, keywords=None, defaults=(None,)), ('document', 'c7e4cfffc93ae89c8f6f53b6d650f923')) +paddle.fluid.layers.diag (ArgSpec(args=['diagonal'], varargs=None, keywords=None, defaults=None), ('document', '2964d07340e32e47efb6e5db619875c7')) paddle.fluid.layers.While.__init__ (ArgSpec(args=['self', 'cond', 'is_test', 'name'], varargs=None, keywords=None, defaults=(False, None)), ('document', '6adf97f83acf6453d4a6a4b1070f3754')) paddle.fluid.layers.While.block (ArgSpec(args=['self'], varargs=None, keywords=None, defaults=None), ('document', '6adf97f83acf6453d4a6a4b1070f3754')) paddle.fluid.layers.Switch.__init__ (ArgSpec(args=['self', 'name'], varargs=None, keywords=None, defaults=(None,)), ('document', '6adf97f83acf6453d4a6a4b1070f3754')) diff --git a/paddle/fluid/operators/diag_op.cc b/paddle/fluid/operators/diag_op.cc new file mode 100644 index 0000000000..5fb18a1d69 --- /dev/null +++ b/paddle/fluid/operators/diag_op.cc @@ -0,0 +1,60 @@ +/* Copyright (c) 2019 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/diag_op.h" + +namespace paddle { +namespace operators { + +class DiagOp : public framework::OperatorWithKernel { + public: + using framework::OperatorWithKernel::OperatorWithKernel; + + void InferShape(framework::InferShapeContext *ctx) const override { + PADDLE_ENFORCE(ctx->HasInput("Diagonal"), + "Input(Diagonal) of DiagOp should not be null."); + + PADDLE_ENFORCE(ctx->HasOutput("Out"), + "Output(Out) of DiagOp should not be null."); + + auto s_dims = ctx->GetInputDim("Diagonal"); + PADDLE_ENFORCE(s_dims.size() == 1, + "The rank of Input(Diagonal) should only be 1."); + + ctx->SetOutputDim("Out", {s_dims[0], s_dims[0]}); + } +}; + +class DiagOpMaker : public framework::OpProtoAndCheckerMaker { + public: + void Make() override { + AddInput("Diagonal", + "Diagonal values of square matrix. It is a tensor with rank 1."); + AddOutput("Out", "A square matrix."); + AddComment(R"DOC( + Return a square matrix with specified diagonal values. +)DOC"); + } +}; +} // namespace operators +} // namespace paddle + +namespace ops = paddle::operators; +REGISTER_OPERATOR(diag, ops::DiagOp, ops::DiagOpMaker, + paddle::framework::EmptyGradOpMaker); +REGISTER_OP_CPU_KERNEL( + diag, ops::DiagKernel, + ops::DiagKernel, + ops::DiagKernel, + ops::DiagKernel); diff --git a/paddle/fluid/operators/diag_op.cu b/paddle/fluid/operators/diag_op.cu new file mode 100644 index 0000000000..9fe1b83b66 --- /dev/null +++ b/paddle/fluid/operators/diag_op.cu @@ -0,0 +1,23 @@ +/* Copyright (c) 2019 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/framework/op_registry.h" +#include "paddle/fluid/operators/diag_op.h" + +namespace ops = paddle::operators; +REGISTER_OP_CUDA_KERNEL( + diag, ops::DiagKernel, + ops::DiagKernel, + ops::DiagKernel, + ops::DiagKernel); diff --git a/paddle/fluid/operators/diag_op.h b/paddle/fluid/operators/diag_op.h new file mode 100644 index 0000000000..f89415ae08 --- /dev/null +++ b/paddle/fluid/operators/diag_op.h @@ -0,0 +1,59 @@ +/* Copyright (c) 2019 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 "paddle/fluid/framework/op_registry.h" +#include "paddle/fluid/operators/math/math_function.h" +#include "paddle/fluid/platform/for_range.h" + +namespace paddle { +namespace operators { + +template +struct DiagFunctor { + DiagFunctor(const T* diagonal, int64_t numel, T* output) + : diagonal_(diagonal), numel_(numel), output_(output) {} + + HOSTDEVICE void operator()(size_t idx) const { + output_[idx * numel_ + idx] = diagonal_[idx]; + } + + const T* diagonal_; + int64_t numel_; + T* output_; +}; + +template +class DiagKernel : public framework::OpKernel { + public: + void Compute(const framework::ExecutionContext& context) const override { + auto* diagonal = context.Input("Diagonal"); + auto* diag_data = diagonal->data(); + auto numel = diagonal->numel(); + auto* out = context.Output("Out"); + T* out_data = out->mutable_data(context.GetPlace()); + + math::SetConstant set_zero; + auto& dev_ctx = context.template device_context(); + set_zero(dev_ctx, out, static_cast(0)); + + platform::ForRange for_range(dev_ctx, numel); + DiagFunctor functor(diag_data, numel, out_data); + for_range(functor); + } +}; + +} // namespace operators +} // namespace paddle diff --git a/python/paddle/fluid/layers/tensor.py b/python/paddle/fluid/layers/tensor.py index 45856670e2..9a0afcd451 100644 --- a/python/paddle/fluid/layers/tensor.py +++ b/python/paddle/fluid/layers/tensor.py @@ -28,7 +28,7 @@ __all__ = [ 'tensor_array_to_tensor', 'concat', 'sums', 'assign', 'fill_constant_batch_size_like', 'fill_constant', 'argmin', 'argmax', 'argsort', 'ones', 'zeros', 'reverse', 'has_inf', 'has_nan', 'isfinite', - 'range', 'linspace', 'zeros_like' + 'range', 'linspace', 'zeros_like', 'diag' ] @@ -890,3 +890,39 @@ def zeros_like(x, out=None): type='fill_zeros_like', inputs={'X': [x]}, outputs={'Out': [out]}) out.stop_gradient = True return out + + +def diag(diagonal): + """ + **diag** + + This function creates a square matrix which has diagonal values specified by `diagonal`. + + Args: + diagonal(Variable|numpy.ndarray): The input tensor specifying diagonal values, should be of rank 1. + + Returns: + Variable: The tensor variable storing the square matrix. + + Examples: + .. code-block:: python + + # [[3, 0, 0] + # [0, 4, 0] + # [0, 0, 5] + data = fluid.layers.diag(np.arange(3, 6)) + + """ + + helper = LayerHelper("diag", **locals()) + + if not isinstance(diagonal, Variable): + diagonal = assign(diagonal) + + out = helper.create_variable_for_type_inference(dtype=diagonal.dtype) + + helper.append_op( + type='diag', inputs={'Diagonal': [diagonal]}, outputs={'Out': [out]}) + + out.stop_gradient = True + return out diff --git a/python/paddle/fluid/tests/unittests/test_diag.py b/python/paddle/fluid/tests/unittests/test_diag.py new file mode 100644 index 0000000000..eed8b91f0e --- /dev/null +++ b/python/paddle/fluid/tests/unittests/test_diag.py @@ -0,0 +1,43 @@ +# Copyright (c) 2019 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 TestDiagOp(OpTest): + def setUp(self): + self.op_type = "diag" + self.init_config() + self.inputs = {'Diagonal': self.case} + + self.outputs = {'Out': np.diag(self.inputs['Diagonal'])} + + def test_check_output(self): + self.check_output() + + def init_config(self): + self.case = np.arange(3, 6) + + +class TestDiagOpCase1(TestDiagOp): + def init_config(self): + self.case = np.array([3], dtype='int32') + + +if __name__ == "__main__": + unittest.main() -- GitLab