From 217db273371abd7b78c4a777992a6090c7e4d0ba Mon Sep 17 00:00:00 2001 From: dengkaipeng Date: Tue, 5 Mar 2019 03:55:33 +0000 Subject: [PATCH] add mkldnn support. test=develop --- .../operators/mkldnn/softmax_mkldnn_op.cc | 128 +++++++++++++----- paddle/fluid/operators/softmax_cudnn_op.cu.cc | 1 - paddle/fluid/operators/softmax_op.cc | 11 +- python/paddle/fluid/layers/nn.py | 17 ++- .../fluid/tests/unittests/test_layers.py | 2 +- 5 files changed, 111 insertions(+), 48 deletions(-) diff --git a/paddle/fluid/operators/mkldnn/softmax_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/softmax_mkldnn_op.cc index 0ce55221945..4e4f482987a 100644 --- a/paddle/fluid/operators/mkldnn/softmax_mkldnn_op.cc +++ b/paddle/fluid/operators/mkldnn/softmax_mkldnn_op.cc @@ -110,28 +110,51 @@ class SoftmaxMKLDNNKernel : public paddle::framework::OpKernel { "It must use CPUPlace."); auto& dev_ctx = ctx.template device_context(); auto mkldnn_engine = dev_ctx.GetEngine(); - const Tensor* input = ctx.Input("X"); - Tensor* output = ctx.Output("Out"); + const Tensor* X = ctx.Input("X"); + Tensor* Out = ctx.Output("Out"); PADDLE_ENFORCE_EQ( - input->dims(), output->dims(), + X->dims(), Out->dims(), "The shape of softmax's input and output must be identical."); + const int axis = ctx.Attr("axis"); + int rank = X->dims().size(); + // make sure 'output' holds memory, which will be shared by // 'flattened_output' later. - output->mutable_data(ctx.GetPlace()); + Out->mutable_data(ctx.GetPlace()); + + std::vector perm, shape; + CalcTransPermAndShapeByAxis(*X, axis, &perm, &shape); + + Tensor X_2d, Out_2d; + Tensor X_trans, Out_trans; + if (axis != -1 && axis != rank - 1) { + X_trans.mutable_data(framework::make_ddim(shape), ctx.GetPlace()); + Out_trans.mutable_data(framework::make_ddim(shape), ctx.GetPlace()); + TransCompute(rank, dev_ctx, *X, &X_trans, perm); + TransCompute(rank, dev_ctx, *Out, &Out_trans, perm); + X_2d = framework::ReshapeToMatrix(X_trans, rank - 1); + Out_2d = framework::ReshapeToMatrix(Out_trans, rank - 1); + } else { + X_2d = framework::ReshapeToMatrix(*X, rank - 1); + Out_2d = framework::ReshapeToMatrix(*Out, rank - 1); + } // flatten input and output to 2-D matrixs - auto dims = input->dims(); // input and output share the same shape - auto flattened_dims = framework::flatten_to_2d(dims, dims.size() - 1); - framework::Tensor flattened_input; - framework::Tensor flattened_output; - flattened_input.ShareDataWith(*input).Resize(flattened_dims); - flattened_output.ShareDataWith(*output).Resize(flattened_dims); - - const T* input_data = flattened_input.data(); - T* output_data = flattened_output.mutable_data(ctx.GetPlace()); - - std::vector src_tz = paddle::framework::vectorize2int(flattened_dims); + // auto dims = input->dims(); // input and output share the same shape + // auto flattened_dims = framework::flatten_to_2d(dims, dims.size() - 1); + // framework::Tensor flattened_input; + // framework::Tensor flattened_output; + // flattened_input.ShareDataWith(*input).Resize(flattened_dims); + // flattened_output.ShareDataWith(*output).Resize(flattened_dims); + + // const T* input_data = flattened_input.data(); + // T* output_data = flattened_output.mutable_data(ctx.GetPlace()); + const T* input_data = X_2d.data(); + T* output_data = Out_2d.mutable_data(ctx.GetPlace()); + + // std::vector src_tz = paddle::framework::vectorize2int(flattened_dims); + std::vector src_tz = paddle::framework::vectorize2int(X_2d.dims()); std::vector dst_tz = src_tz; // Same memory descriptor to be used for input and output memory::dims softmax_tz = {src_tz[0], src_tz[1]}; @@ -178,6 +201,10 @@ class SoftmaxMKLDNNKernel : public paddle::framework::OpKernel { output_data[i] < threshold ? threshold : output_data[i]; } } + + if (axis != -1 && axis != rank - 1) { + TransCompute(rank, dev_ctx, Out_trans, Out, perm); + } } }; @@ -190,33 +217,60 @@ class SoftmaxMKLDNNGradKernel : public paddle::framework::OpKernel { auto& dev_ctx = ctx.template device_context(); auto mkldnn_engine = dev_ctx.GetEngine(); - const Tensor* output = ctx.Input("Out"); - auto* dout = ctx.template Input(framework::GradVarName("Out")); - auto* dx = + const Tensor* Out = ctx.Input("Out"); + auto* dOut = ctx.template Input(framework::GradVarName("Out")); + auto* dX = ctx.template Output(framework::GradVarName("X")); PADDLE_ENFORCE_EQ( - dout->dims(), dx->dims(), + dOut->dims(), dX->dims(), "The shape of softmax_grad's input and output must be identical."); + const int axis = ctx.Attr("axis"); + int rank = Out->dims().size(); + // make sure 'dx' holds memory, which will be shared by 'flattened_dx' // later. - dx->template mutable_data(ctx.GetPlace()); - - auto dims = dout->dims(); // input and output share the same shape - auto flattened_dims = framework::flatten_to_2d(dims, dims.size() - 1); - framework::Tensor flattened_output; - framework::Tensor flattened_dout; - framework::Tensor flattened_dx; - flattened_output.ShareDataWith(*output).Resize(flattened_dims); - flattened_dout.ShareDataWith(*dout).Resize(flattened_dims); - flattened_dx.ShareDataWith(*dx).Resize(flattened_dims); - - const T* dst_data = flattened_output.data(); - const T* diff_dst_ptr = flattened_dout.template data(); - T* diff_src_ptr = flattened_dx.template mutable_data(ctx.GetPlace()); - - std::vector dst_tz = paddle::framework::vectorize2int(flattened_dims); + dX->template mutable_data(ctx.GetPlace()); + + std::vector perm, shape; + CalcTransPermAndShapeByAxis(*dX, axis, &perm, &shape); + + Tensor dX_2d, Out_2d, dOut_2d; + Tensor dX_trans, Out_trans, dOut_trans; + if (axis != -1 && axis != rank - 1) { + dX_trans.mutable_data(framework::make_ddim(shape), ctx.GetPlace()); + Out_trans.mutable_data(framework::make_ddim(shape), ctx.GetPlace()); + dOut_trans.mutable_data(framework::make_ddim(shape), ctx.GetPlace()); + TransCompute(rank, dev_ctx, *dX, &dX_trans, perm); + TransCompute(rank, dev_ctx, *Out, &Out_trans, perm); + TransCompute(rank, dev_ctx, *dOut, &dOut_trans, perm); + dX_2d = framework::ReshapeToMatrix(dX_trans, rank - 1); + Out_2d = framework::ReshapeToMatrix(Out_trans, rank - 1); + dOut_2d = framework::ReshapeToMatrix(dOut_trans, rank - 1); + } else { + dX_2d = framework::ReshapeToMatrix(*dX, rank - 1); + Out_2d = framework::ReshapeToMatrix(*Out, rank - 1); + dOut_2d = framework::ReshapeToMatrix(*dOut, rank - 1); + } + + // auto dims = dout->dims(); // input and output share the same shape + // auto flattened_dims = framework::flatten_to_2d(dims, dims.size() - 1); + // framework::Tensor flattened_output; + // framework::Tensor flattened_dout; + // framework::Tensor flattened_dx; + // flattened_output.ShareDataWith(*output).Resize(flattened_dims); + // flattened_dout.ShareDataWith(*dout).Resize(flattened_dims); + // flattened_dx.ShareDataWith(*dx).Resize(flattened_dims); + + // const T* dst_data = flattened_output.data(); + // const T* diff_dst_ptr = flattened_dout.template data(); + // T* diff_src_ptr = flattened_dx.template mutable_data(ctx.GetPlace()); + const T* dst_data = Out_2d.data(); + const T* diff_dst_ptr = dOut_2d.template data(); + T* diff_src_ptr = dX_2d.template mutable_data(ctx.GetPlace()); + + std::vector dst_tz = paddle::framework::vectorize2int(Out_2d.dims()); std::vector src_tz(dst_tz); // Same memory descriptor to be used for input and output @@ -261,6 +315,10 @@ class SoftmaxMKLDNNGradKernel : public paddle::framework::OpKernel { std::vector pipeline{*softmax_bwd_p}; stream(stream::kind::eager).submit(pipeline).wait(); + + if (axis != -1 && axis != rank - 1) { + TransCompute(rank, dev_ctx, dX_trans, dX, perm); + } } }; } // namespace operators diff --git a/paddle/fluid/operators/softmax_cudnn_op.cu.cc b/paddle/fluid/operators/softmax_cudnn_op.cu.cc index 84151d70b99..dc5b7bb0af4 100644 --- a/paddle/fluid/operators/softmax_cudnn_op.cu.cc +++ b/paddle/fluid/operators/softmax_cudnn_op.cu.cc @@ -28,7 +28,6 @@ class SoftmaxCUDNNKernel : public framework::OpKernel { auto& dev_ctx = context.template device_context(); auto* X = context.Input("X"); auto* Out = context.Output("Out"); - // auto dims = X->dims(); const int axis = context.Attr("axis"); int rank = X->dims().size(); diff --git a/paddle/fluid/operators/softmax_op.cc b/paddle/fluid/operators/softmax_op.cc index bd3b14775fc..02f256fa644 100644 --- a/paddle/fluid/operators/softmax_op.cc +++ b/paddle/fluid/operators/softmax_op.cc @@ -85,10 +85,10 @@ class SoftmaxOpMaker : public framework::OpProtoAndCheckerMaker { void Make() override { AddInput("X", "The input tensor of softmax, " - "whose last dimension is the input_feature_dimensions."); + "whose :attr:`axis` dimension is the input_feature_dimensions."); AddOutput("Out", "The normalized values with the same shape as X."); AddAttr("axis", - "The dimension of Input(x) to perform softmax," + "The dimension index of Input(x) to perform softmax," "default -1 for last dimension") .SetDefault(-1); AddAttr( @@ -115,12 +115,13 @@ Softmax Operator. The input of the softmax operator is a tensor of any rank. The output tensor has the same shape as the input. -The input tensor will first be logically flattened to a 2-D matrix. The matrix's -second dimension(row length) is as same as the last dimension of the input +The :attr:`axis` th dimension of the input tensor will be permuted to the last. +Then the input tensor will be logically flattened to a 2-D matrix. The matrix's +second dimension(row length) is as same as the :attr:`axis` dimension of the input tensor, and the first dimension(column length) is the product of all other dimensions of the input tensor. For each row of the matrix, the softmax operator squashes the K-dimensional(K is the width of the matrix, which is also the size -of the input tensor's last dimension) vector of arbitrary real values to a +of the input tensor's :attr:`axis` dimension) vector of arbitrary real values to a K-dimensional vector of real values in the range [0, 1] that add up to 1. It computes the exponential of the given dimension and the sum of exponential values of all the other dimensions in the K-dimensional vector input. diff --git a/python/paddle/fluid/layers/nn.py b/python/paddle/fluid/layers/nn.py index dbe495b75c8..273d74ca6ed 100644 --- a/python/paddle/fluid/layers/nn.py +++ b/python/paddle/fluid/layers/nn.py @@ -1819,17 +1819,18 @@ def sequence_softmax(input, use_cudnn=False, name=None): return softmax_out -def softmax(input, use_cudnn=False, name=None): +def softmax(input, use_cudnn=False, name=None, axis=-1): """ The input of the softmax operator is a tensor of any rank. The output tensor has the same shape as the input. - The input tensor will first be logically flattened to a 2-D matrix. The matrix's - second dimension(row length) is as same as the last dimension of the input + The :attr:`axis` th dimension of the input tensor will be permuted to the last. + Then the input tensor will be logically flattened to a 2-D matrix. The matrix's + second dimension(row length) is as same as the :attr:`axis` th dimension of the input tensor, and the first dimension(column length) is the product of all other dimensions of the input tensor. For each row of the matrix, the softmax operator squashes the K-dimensional(K is the width of the matrix, which is also the size - of the input tensor's last dimension) vector of arbitrary real values to a + of the input tensor's :attr:`axis` th dimension) vector of arbitrary real values to a K-dimensional vector of real values in the range [0, 1] that add up to 1. It computes the exponential of the given dimension and the sum of exponential @@ -1851,6 +1852,7 @@ def softmax(input, use_cudnn=False, name=None): False by default. Default: False name (str|None): A name for this layer(optional). If set None, the layer will be named automatically. Default: None. + axis (int): The index of dimension to perform softmax calculation. Default: -1. Returns: Variable: output of softmax @@ -1860,7 +1862,7 @@ def softmax(input, use_cudnn=False, name=None): .. code-block:: python fc = fluid.layers.fc(input=x, size=10) - softmax = fluid.layers.softmax(input=fc) + softmax = fluid.layers.softmax(input=fc, axis=1) """ helper = LayerHelper('softmax', **locals()) @@ -1870,7 +1872,10 @@ def softmax(input, use_cudnn=False, name=None): type="softmax", inputs={"X": input}, outputs={"Out": softmax_out}, - attrs={"use_cudnn": use_cudnn}) + attrs={ + "axis": axis, + "use_cudnn": use_cudnn + }) return softmax_out diff --git a/python/paddle/fluid/tests/unittests/test_layers.py b/python/paddle/fluid/tests/unittests/test_layers.py index 885ee170e80..4e255293b66 100644 --- a/python/paddle/fluid/tests/unittests/test_layers.py +++ b/python/paddle/fluid/tests/unittests/test_layers.py @@ -513,7 +513,7 @@ class TestBook(unittest.TestCase): with program_guard(program): data = layers.data(name='data', shape=[10], dtype='float32') hid = layers.fc(input=data, size=20) - self.assertIsNotNone(layers.softmax(hid)) + self.assertIsNotNone(layers.softmax(hid, axis=1)) print(str(program)) def test_space_to_depth(self): -- GitLab