Add trace op (#23873)
* add trace op, test=develop * Optimized the implementation of trace op, test=develop * fix a bug of include in trace_op.h, test=develop * move trace API from creation to math, test=develop * modified en doc. test=develop * add complex trace api * add complex sum api, test=develop * modified en doc of complex sum and trace, test=develop * modified doc and trace API, test=develop * modified en doc of trace and sum, test=develop * modified comment in complex kron API, test=develop * OP Should Not Have Unused Input, test=develop * add GetExpectedKernelType, test=developrevert-24314-dev/fix_err_msg
Li Fuchen
5 years ago
committed by
GitHub
parent
fa43d74a3a
commit
077e5a0fe5
@ -0,0 +1,172 @@
|
||||
// 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/trace_op.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace operators {
|
||||
|
||||
class TraceOp : public framework::OperatorWithKernel {
|
||||
public:
|
||||
using framework::OperatorWithKernel::OperatorWithKernel;
|
||||
|
||||
void InferShape(framework::InferShapeContext *ctx) const override {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
ctx->HasInput("Input"), true,
|
||||
platform::errors::NotFound("Input of TraceOp is not found."));
|
||||
|
||||
PADDLE_ENFORCE_EQ(
|
||||
ctx->HasOutput("Out"), true,
|
||||
platform::errors::NotFound("Output of TraceOp is not found."));
|
||||
|
||||
int dim1 = ctx->Attrs().Get<int>("dim1");
|
||||
int dim2 = ctx->Attrs().Get<int>("dim2");
|
||||
|
||||
auto x_dims = ctx->GetInputDim("Input");
|
||||
|
||||
int dim1_ = dim1 < 0 ? x_dims.size() + dim1 : dim1;
|
||||
int dim2_ = dim2 < 0 ? x_dims.size() + dim2 : dim2;
|
||||
|
||||
PADDLE_ENFORCE_GE(
|
||||
x_dims.size(), 2,
|
||||
platform::errors::OutOfRange(
|
||||
"trace requires an tensor of at least two dimensions"));
|
||||
PADDLE_ENFORCE_LT(
|
||||
dim1_, x_dims.size(),
|
||||
platform::errors::OutOfRange(
|
||||
"Attr(dim1) is out of range (expected to be in range of [%ld, "
|
||||
"%ld], but got %ld).",
|
||||
-(x_dims.size()), (x_dims.size() - 1), dim1));
|
||||
PADDLE_ENFORCE_LT(
|
||||
dim2_, x_dims.size(),
|
||||
platform::errors::OutOfRange(
|
||||
"Attr(dim2) is out of range (expected to be in range of [%ld, "
|
||||
"%ld], but got %ld).",
|
||||
-(x_dims.size()), (x_dims.size() - 1), dim2));
|
||||
PADDLE_ENFORCE_NE(dim1_, dim2_,
|
||||
platform::errors::InvalidArgument(
|
||||
"The dimensions should not be identical "
|
||||
"%ld vs %ld.",
|
||||
dim1, dim2));
|
||||
|
||||
auto sizes = vectorize(x_dims);
|
||||
if (x_dims.size() == 2) {
|
||||
sizes.clear();
|
||||
sizes.push_back(1);
|
||||
} else {
|
||||
sizes.erase(sizes.begin() + std::max(dim1_, dim2_));
|
||||
sizes.erase(sizes.begin() + std::min(dim1_, dim2_));
|
||||
}
|
||||
ctx->SetOutputDim("Out", framework::make_ddim(sizes));
|
||||
}
|
||||
};
|
||||
|
||||
class TraceOpMaker : public framework::OpProtoAndCheckerMaker {
|
||||
public:
|
||||
void Make() override {
|
||||
AddInput("Input",
|
||||
"(Tensor) The input tensor, from which the diagonals are taken.");
|
||||
AddOutput("Out", "(Tensor) the sum along diagonals of the input tensor");
|
||||
AddAttr<int>(
|
||||
"offset",
|
||||
R"DOC((int, default 0), offset of the diagonal from the main diagonal. Can be both positive and negative. Defaults to 0.
|
||||
)DOC")
|
||||
.SetDefault(0);
|
||||
AddAttr<int>(
|
||||
"dim1",
|
||||
R"DOC((int, default 0), the first dim of the 2-D planes from which the diagonals should be taken.
|
||||
Can be both positive and negative. Default: 0.
|
||||
)DOC")
|
||||
.SetDefault(-2);
|
||||
AddAttr<int>(
|
||||
"dim2",
|
||||
R"DOC((int, default 1), the second dim of the 2-D planes from which the diagonals should be taken.
|
||||
Can be both positive and negative. Default: 1.
|
||||
)DOC")
|
||||
.SetDefault(-1);
|
||||
AddComment(R"DOC(
|
||||
Trace Operator.
|
||||
Return the sum along diagonals of the input tensor.
|
||||
The behavior of this operator is similar to how `numpy.trace` works.
|
||||
|
||||
If Input is 2-D, returns the sum of diagonal.
|
||||
If Input has larger dimensions, then returns an tensor of diagonals sum, diagonals be taken from
|
||||
the 2-D planes specified by dim1 and dim2.
|
||||
|
||||
)DOC");
|
||||
}
|
||||
};
|
||||
class TraceOpGrad : public framework::OperatorWithKernel {
|
||||
public:
|
||||
using framework::OperatorWithKernel::OperatorWithKernel;
|
||||
|
||||
void InferShape(framework::InferShapeContext *ctx) const override {
|
||||
PADDLE_ENFORCE_EQ(
|
||||
ctx->HasInput("Input"), true,
|
||||
platform::errors::NotFound("Input(Input) of TraceOp is not found."));
|
||||
PADDLE_ENFORCE_EQ(ctx->HasOutput(framework::GradVarName("Input")), true,
|
||||
platform::errors::NotFound(
|
||||
"Output(Input@GRAD) of TraceGradOp is not found."));
|
||||
ctx->SetOutputDim(framework::GradVarName("Input"),
|
||||
ctx->GetInputDim("Input"));
|
||||
}
|
||||
|
||||
protected:
|
||||
framework::OpKernelType GetExpectedKernelType(
|
||||
const framework::ExecutionContext &ctx) const override {
|
||||
return framework::OpKernelType(OperatorWithKernel::IndicateVarDataType(
|
||||
ctx, framework::GradVarName("Out")),
|
||||
ctx.GetPlace());
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
class TraceGradOpMaker : public framework::SingleGradOpMaker<T> {
|
||||
public:
|
||||
using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
|
||||
|
||||
protected:
|
||||
void Apply(GradOpPtr<T> grad_op) const override {
|
||||
grad_op->SetType("trace_grad");
|
||||
grad_op->SetInput("Input", this->Input("Input"));
|
||||
grad_op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
|
||||
grad_op->SetOutput(framework::GradVarName("Input"),
|
||||
this->InputGrad("Input"));
|
||||
grad_op->SetAttrMap(this->Attrs());
|
||||
}
|
||||
};
|
||||
|
||||
DECLARE_NO_NEED_BUFFER_VARS_INFERER(TraceGradNoNeedBufferVarsInference,
|
||||
"Input");
|
||||
|
||||
} // namespace operators
|
||||
} // namespace paddle
|
||||
|
||||
namespace ops = paddle::operators;
|
||||
REGISTER_OPERATOR(trace, ops::TraceOp, ops::TraceOpMaker,
|
||||
ops::TraceGradOpMaker<paddle::framework::OpDesc>,
|
||||
ops::TraceGradOpMaker<paddle::imperative::OpBase>);
|
||||
|
||||
REGISTER_OPERATOR(trace_grad, ops::TraceOpGrad,
|
||||
ops::TraceGradNoNeedBufferVarsInference);
|
||||
REGISTER_OP_CPU_KERNEL(
|
||||
trace, ops::TraceKernel<paddle::platform::CPUDeviceContext, int>,
|
||||
ops::TraceKernel<paddle::platform::CPUDeviceContext, float>,
|
||||
ops::TraceKernel<paddle::platform::CPUDeviceContext, double>,
|
||||
ops::TraceKernel<paddle::platform::CPUDeviceContext, int64_t>);
|
||||
REGISTER_OP_CPU_KERNEL(
|
||||
trace_grad, ops::TraceGradKernel<paddle::platform::CPUDeviceContext, int>,
|
||||
ops::TraceGradKernel<paddle::platform::CPUDeviceContext, float>,
|
||||
ops::TraceGradKernel<paddle::platform::CPUDeviceContext, double>,
|
||||
ops::TraceGradKernel<paddle::platform::CPUDeviceContext, int64_t>);
|
||||
@ -0,0 +1,70 @@
|
||||
// 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/reduce_ops/cub_reduce.h"
|
||||
#include "paddle/fluid/operators/trace_op.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace operators {
|
||||
|
||||
template <typename T>
|
||||
struct IdentityFunctor {
|
||||
HOSTDEVICE explicit inline IdentityFunctor() {}
|
||||
|
||||
HOSTDEVICE inline T operator()(const T& x) const { return x; }
|
||||
};
|
||||
|
||||
template <typename DeviceContext, typename T>
|
||||
class TraceCUDAKernel : public framework::OpKernel<T> {
|
||||
public:
|
||||
void Compute(const framework::ExecutionContext& context) const override {
|
||||
auto* input = context.Input<framework::Tensor>("Input");
|
||||
auto* out = context.Output<framework::Tensor>("Out");
|
||||
|
||||
const int64_t offset = context.Attr<int>("offset");
|
||||
const int64_t dim1 = context.Attr<int>("dim1");
|
||||
const int64_t dim2 = context.Attr<int>("dim2");
|
||||
|
||||
T* out_data = out->mutable_data<T>(context.GetPlace());
|
||||
const framework::Tensor diag =
|
||||
Diagonal<DeviceContext, T>(context, input, offset, dim1, dim2);
|
||||
if (diag.numel() > 0) {
|
||||
auto stream = context.cuda_device_context().stream();
|
||||
std::vector<int> reduce_dims;
|
||||
reduce_dims.push_back(out->dims().size());
|
||||
TensorReduce<T, T, cub::Sum, IdentityFunctor<T>>(
|
||||
diag, out, reduce_dims, static_cast<T>(0), cub::Sum(),
|
||||
IdentityFunctor<T>(), stream);
|
||||
}
|
||||
}
|
||||
};
|
||||
} // namespace operators
|
||||
} // namespace paddle
|
||||
|
||||
namespace ops = paddle::operators;
|
||||
namespace platform = paddle::platform;
|
||||
REGISTER_OP_CUDA_KERNEL(
|
||||
trace, ops::TraceCUDAKernel<paddle::platform::CUDADeviceContext, int>,
|
||||
ops::TraceCUDAKernel<paddle::platform::CUDADeviceContext, int64_t>,
|
||||
ops::TraceCUDAKernel<paddle::platform::CUDADeviceContext,
|
||||
platform::float16>,
|
||||
ops::TraceCUDAKernel<paddle::platform::CUDADeviceContext, float>,
|
||||
ops::TraceCUDAKernel<paddle::platform::CUDADeviceContext, double>);
|
||||
REGISTER_OP_CUDA_KERNEL(
|
||||
trace_grad, ops::TraceGradKernel<paddle::platform::CUDADeviceContext, int>,
|
||||
ops::TraceGradKernel<paddle::platform::CUDADeviceContext, int64_t>,
|
||||
ops::TraceGradKernel<paddle::platform::CUDADeviceContext,
|
||||
platform::float16>,
|
||||
ops::TraceGradKernel<paddle::platform::CUDADeviceContext, float>,
|
||||
ops::TraceGradKernel<paddle::platform::CUDADeviceContext, double>);
|
||||
File diff suppressed because it is too large
Load Diff
@ -0,0 +1,42 @@
|
||||
# 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.
|
||||
|
||||
import unittest
|
||||
import numpy as np
|
||||
from numpy.random import random as rand
|
||||
import paddle.complex as cpx
|
||||
import paddle.fluid as fluid
|
||||
import paddle.fluid.dygraph as dg
|
||||
|
||||
|
||||
class TestComplexSumLayer(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._dtype = "float64"
|
||||
self._places = [fluid.CPUPlace()]
|
||||
if fluid.core.is_compiled_with_cuda():
|
||||
self._places.append(fluid.CUDAPlace(0))
|
||||
|
||||
def test_complex_x(self):
|
||||
input = rand([2, 10, 10]).astype(self._dtype) + 1j * rand(
|
||||
[2, 10, 10]).astype(self._dtype)
|
||||
for place in self._places:
|
||||
with dg.guard(place):
|
||||
var_x = dg.to_variable(input)
|
||||
result = cpx.sum(var_x, dim=[1, 2]).numpy()
|
||||
target = np.sum(input, axis=(1, 2))
|
||||
self.assertTrue(np.allclose(result, target))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@ -0,0 +1,42 @@
|
||||
# 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.
|
||||
|
||||
import unittest
|
||||
import numpy as np
|
||||
from numpy.random import random as rand
|
||||
import paddle.complex as cpx
|
||||
import paddle.fluid as fluid
|
||||
import paddle.fluid.dygraph as dg
|
||||
|
||||
|
||||
class TestComplexTraceLayer(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self._dtype = "float64"
|
||||
self._places = [fluid.CPUPlace()]
|
||||
if fluid.core.is_compiled_with_cuda():
|
||||
self._places.append(fluid.CUDAPlace(0))
|
||||
|
||||
def test_complex_x(self):
|
||||
input = rand([2, 20, 2, 3]).astype(self._dtype) + 1j * rand(
|
||||
[2, 20, 2, 3]).astype(self._dtype)
|
||||
for place in self._places:
|
||||
with dg.guard(place):
|
||||
var_x = dg.to_variable(input)
|
||||
result = cpx.trace(var_x, offset=1, dim1=0, dim2=2).numpy()
|
||||
target = np.trace(input, offset=1, axis1=0, axis2=2)
|
||||
self.assertTrue(np.allclose(result, target))
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
@ -0,0 +1,89 @@
|
||||
# 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
|
||||
import paddle.nn.functional as F
|
||||
import paddle.fluid as fluid
|
||||
import paddle.fluid.core as core
|
||||
import paddle.tensor as tensor
|
||||
|
||||
|
||||
class TestTraceOp(OpTest):
|
||||
def setUp(self):
|
||||
self.op_type = "trace"
|
||||
self.init_config()
|
||||
self.outputs = {'Out': self.target}
|
||||
|
||||
def test_check_output(self):
|
||||
self.check_output()
|
||||
|
||||
def test_check_grad(self):
|
||||
self.check_grad(['Input'], 'Out')
|
||||
|
||||
def init_config(self):
|
||||
self.case = np.random.randn(20, 6).astype('float64')
|
||||
self.inputs = {'Input': self.case}
|
||||
self.attrs = {'offset': 0, 'dim1': 0, 'dim2': 1}
|
||||
self.target = np.trace(self.inputs['Input'])
|
||||
|
||||
|
||||
class TestTraceOpCase1(TestTraceOp):
|
||||
def init_config(self):
|
||||
self.case = np.random.randn(2, 20, 2, 3).astype('float32')
|
||||
self.inputs = {'Input': self.case}
|
||||
self.attrs = {'offset': 1, 'dim1': 0, 'dim2': 2}
|
||||
self.target = np.trace(
|
||||
self.inputs['Input'],
|
||||
offset=self.attrs['offset'],
|
||||
axis1=self.attrs['dim1'],
|
||||
axis2=self.attrs['dim2'])
|
||||
|
||||
|
||||
class TestTraceOpCase2(TestTraceOp):
|
||||
def init_config(self):
|
||||
self.case = np.random.randn(2, 20, 2, 3).astype('float32')
|
||||
self.inputs = {'Input': self.case}
|
||||
self.attrs = {'offset': -5, 'dim1': 1, 'dim2': -1}
|
||||
self.target = np.trace(
|
||||
self.inputs['Input'],
|
||||
offset=self.attrs['offset'],
|
||||
axis1=self.attrs['dim1'],
|
||||
axis2=self.attrs['dim2'])
|
||||
|
||||
|
||||
class TestTraceAPICase(unittest.TestCase):
|
||||
def test_case1(self):
|
||||
case = np.random.randn(2, 20, 2, 3).astype('float32')
|
||||
data1 = fluid.data(name='data1', shape=[2, 20, 2, 3], dtype='float32')
|
||||
out1 = tensor.trace(data1)
|
||||
out2 = tensor.trace(data1, offset=-5, dim1=1, dim2=-1)
|
||||
|
||||
place = core.CPUPlace()
|
||||
exe = fluid.Executor(place)
|
||||
results = exe.run(fluid.default_main_program(),
|
||||
feed={"data1": case},
|
||||
fetch_list=[out1, out2],
|
||||
return_numpy=True)
|
||||
target1 = np.trace(case)
|
||||
target2 = np.trace(case, offset=-5, axis1=1, axis2=-1)
|
||||
self.assertTrue(np.allclose(results[0], target1))
|
||||
self.assertTrue(np.allclose(results[1], target2))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Loading…
Reference in new issue