diff --git a/paddle/fluid/operators/mkldnn/softmax_mkldnn_op.cc b/paddle/fluid/operators/mkldnn/softmax_mkldnn_op.cc index 0ce552219458859e147ba207c94270bf84a1fe75..4e4f482987a29b8a5a574b503ddf4eb7ea740e3d 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 84151d70b9997ba717d1942b93a9be0c31f534a3..dc5b7bb0af4a484bc97821bf65f12c82340dd273 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 bd3b14775fca49b2c7035186d325db99e54bf293..02f256fa64427a900cb73cc08ef3e737388fbae1 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 dbe495b75c8876de3c3f0fb0abe0357089413254..273d74ca6ed19af945cc83545c42e90ef9656c4c 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 885ee170e8032ef865ebfdd646fed1e995e9e60b..4e255293b6614dee1b4d2173762e6d8a01479889 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):