softmax_op.cc 5.3 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2

Q
Qiao Longfei 已提交
3 4 5
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
6

Q
Qiao Longfei 已提交
7 8 9 10 11 12 13
    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. */
14

Y
Yi Wang 已提交
15
#include "paddle/fluid/operators/softmax_op.h"
K
Kexin Zhao 已提交
16 17 18
#ifdef PADDLE_WITH_CUDA
#include "paddle/fluid/platform/cudnn_helper.h"
#endif
19 20 21 22

namespace paddle {
namespace operators {

D
dongzhihong 已提交
23
class SoftmaxOp : public framework::OperatorWithKernel {
Y
Yu Yang 已提交
24 25 26
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

27
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
28 29
    PADDLE_ENFORCE(ctx->HasInput("X"),
                   "Input(X) of SoftmaxOp should not be null.");
F
fengjiayi 已提交
30 31
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of SoftmaxOp should not be null.");
Q
Qiao Longfei 已提交
32 33 34

    auto x_dims = ctx->GetInputDim("X");
    PADDLE_ENFORCE(x_dims.size() == 2UL,
C
caoying03 已提交
35
                   "The input of softmax op must be a matrix.");
F
fengjiayi 已提交
36
    ctx->SetOutputDim("Out", x_dims);
Q
Qiao Longfei 已提交
37
    ctx->ShareLoD("X", /*->*/ "Out");
38
  }
39 40 41 42 43

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override {
    // choose cudnn kernel if the runtime supported.
K
Kexin Zhao 已提交
44
    framework::LibraryType library_{framework::LibraryType::kPlain};
45
#ifdef PADDLE_WITH_CUDA
K
Kexin Zhao 已提交
46 47
    if (platform::CanCUDNNBeUsed(ctx)) {
      library = framework::LibraryType::kCUDNN;
48 49 50 51 52 53 54
    }
#endif
    std::string data_format = ctx.Attr<std::string>("data_format");
    return framework::OpKernelType(
        framework::ToDataType(ctx.Input<Tensor>("X")->type()), ctx.GetPlace(),
        framework::StringToDataLayout(data_format), library_);
  }
55 56
};

D
dongzhihong 已提交
57
class SoftmaxOpMaker : public framework::OpProtoAndCheckerMaker {
58
 public:
59
  SoftmaxOpMaker(OpProto* proto, OpAttrChecker* op_checker)
60
      : OpProtoAndCheckerMaker(proto, op_checker) {
61
    AddInput("X",
C
caoying03 已提交
62 63
             "The input tensor of softmax. "
             "2-D with shape [batch_size, input_feature_dimensions].");
F
fengjiayi 已提交
64
    AddOutput("Out", "The normalized values with the same shape as X.");
65 66 67 68 69 70 71 72 73 74 75
    AddAttr<bool>(
        "use_cudnn",
        "(bool, default false) Only used in cudnn kernel, need install cudnn")
        .SetDefault(false);
    AddAttr<std::string>(
        "data_format",
        "(string, default NCHW) Only used in "
        "An optional string from: \"NHWC\", \"NCHW\". "
        "Defaults to \"NHWC\". Specify the data format of the output data, "
        "the input will be transformed automatically. ")
        .SetDefault("AnyLayout");
C
caoying03 已提交
76
    AddComment(R"DOC(
77 78 79
Softmax Operator.

The input of the softmax operator is a 2-D tensor with shape N x K (N is the
C
caoying03 已提交
80 81 82 83 84
batch_size, K is the dimension of input feature). The output tensor has the
same shape as the input tensor.

For each row of the input tensor, the softmax operator squashes the
K-dimensional vector of arbitrary real values to a K-dimensional vector of real
85 86 87 88 89 90
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.
Then the ratio of the exponential of the given dimension and the sum of
exponential values of all the other dimensions is the output of the softmax
operator.
C
caoying03 已提交
91

92
For each row $i$ and each column $j$ in Input(X), we have:
F
fengjiayi 已提交
93
    $$Out[i, j] = \frac{\exp(X[i, j])}{\sum_j(exp(X[i, j])}$$
C
caoying03 已提交
94 95

)DOC");
96 97 98
  }
};

D
dongzhihong 已提交
99
class SoftmaxOpGrad : public framework::OperatorWithKernel {
Y
Yu Yang 已提交
100 101 102
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

103
  void InferShape(framework::InferShapeContext* ctx) const override {
F
fengjiayi 已提交
104 105 106 107 108 109
    PADDLE_ENFORCE(ctx->HasInput("Out"), "Input(Out) should be not null.");
    PADDLE_ENFORCE(ctx->HasInput(framework::GradVarName("Out")),
                   "Input(Out@GRAD) should be not null.");
    PADDLE_ENFORCE_EQ(ctx->GetInputDim("Out"),
                      ctx->GetInputDim(framework::GradVarName("Out")),
                      "Input(Out) and its gradients should have a same shape.");
110

Q
Qiao Longfei 已提交
111
    ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("X"));
D
dongzhihong 已提交
112
  }
113 114 115 116 117

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override {
    // choose cudnn kernel if the runtime supported.
K
Kexin Zhao 已提交
118
    framework::LibraryType library_{framework::LibraryType::kPlain};
119
#ifdef PADDLE_WITH_CUDA
K
Kexin Zhao 已提交
120 121
    if (platform::CanCUDNNBeUsed(ctx)) {
      library = framework::LibraryType::kCUDNN;
122 123 124 125 126 127 128
    }
#endif
    std::string data_format = ctx.Attr<std::string>("data_format");
    return framework::OpKernelType(
        framework::ToDataType(ctx.Input<Tensor>("X")->type()), ctx.GetPlace(),
        framework::StringToDataLayout(data_format), library_);
  }
D
dongzhihong 已提交
129 130
};

131 132 133
}  // namespace operators
}  // namespace paddle

D
dongzhihong 已提交
134
namespace ops = paddle::operators;
D
dongzhihong 已提交
135

136 137
REGISTER_OP(softmax, ops::SoftmaxOp, ops::SoftmaxOpMaker, softmax_grad,
            ops::SoftmaxOpGrad);
D
dongzhihong 已提交
138
REGISTER_OP_CPU_KERNEL(
Q
QI JUN 已提交
139 140 141 142
    softmax, ops::SoftmaxKernel<paddle::platform::CPUDeviceContext, float>);
REGISTER_OP_CPU_KERNEL(
    softmax_grad,
    ops::SoftmaxGradKernel<paddle::platform::CPUDeviceContext, float>);