softmax_op.cc 8.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"
16 17 18

#include <string>

K
Kexin Zhao 已提交
19 20 21
#ifdef PADDLE_WITH_CUDA
#include "paddle/fluid/platform/cudnn_helper.h"
#endif
22

23 24 25
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif
26

27 28 29
namespace paddle {
namespace operators {

D
dongzhihong 已提交
30
class SoftmaxOp : public framework::OperatorWithKernel {
Y
Yu Yang 已提交
31 32 33
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

34
  void InferShape(framework::InferShapeContext* ctx) const override {
Q
Qiao Longfei 已提交
35 36
    PADDLE_ENFORCE(ctx->HasInput("X"),
                   "Input(X) of SoftmaxOp should not be null.");
F
fengjiayi 已提交
37 38
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of SoftmaxOp should not be null.");
Q
Qiao Longfei 已提交
39

F
fengjiayi 已提交
40
    ctx->SetOutputDim("Out", ctx->GetInputDim("X"));
Q
Qiao Longfei 已提交
41
    ctx->ShareLoD("X", /*->*/ "Out");
42
  }
43 44 45 46 47

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override {
    // choose cudnn kernel if the runtime supported.
K
Kexin Zhao 已提交
48
    framework::LibraryType library_{framework::LibraryType::kPlain};
M
mozga-intel 已提交
49 50 51
    std::string data_format = ctx.Attr<std::string>("data_format");
    framework::DataLayout layout_ = framework::StringToDataLayout(data_format);

52
#ifdef PADDLE_WITH_CUDA
K
Kexin Zhao 已提交
53
    if (platform::CanCUDNNBeUsed(ctx)) {
K
Kexin Zhao 已提交
54
      library_ = framework::LibraryType::kCUDNN;
55 56
    }
#endif
57 58 59 60
#ifdef PADDLE_WITH_MKLDNN
    if (library_ == framework::LibraryType::kPlain &&
        platform::CanMKLDNNBeUsed(ctx)) {
      library_ = framework::LibraryType::kMKLDNN;
M
mozga-intel 已提交
61
      layout_ = framework::DataLayout::kMKLDNN;
62 63
    }
#endif
K
Kexin Zhao 已提交
64

Y
Yu Yang 已提交
65
    auto input_data_type = ctx.Input<Tensor>("X")->type();
K
Kexin Zhao 已提交
66
    if (input_data_type == framework::proto::VarType::FP16) {
67 68
      PADDLE_ENFORCE(platform::is_gpu_place(ctx.GetPlace()),
                     "float16 can only be used on GPU place");
K
Kexin Zhao 已提交
69 70
    }

M
mozga-intel 已提交
71
    return framework::OpKernelType(input_data_type, ctx.GetPlace(), layout_,
K
Kexin Zhao 已提交
72
                                   library_);
73
  }
74
};
75

D
dongzhihong 已提交
76
class SoftmaxOpMaker : public framework::OpProtoAndCheckerMaker {
77
 public:
Y
Yu Yang 已提交
78
  void Make() override {
79
    AddInput("X",
F
fengjiayi 已提交
80 81
             "The input tensor of softmax, "
             "whose last dimension is the input_feature_dimensions.");
82
    AddOutput("Out", "The normalized values with the same shape as X.");
83 84 85 86 87 88 89 90 91 92 93
    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");
94 95 96
    AddAttr<bool>("use_mkldnn",
                  "(bool, default false) Only used in mkldnn kernel")
        .SetDefault(false);
J
Jacek Czaja 已提交
97
    AddAttr<bool>("is_test",
98 99
                  "(bool, default false) Set to true for inference only, false "
                  "for training. Some layers may run faster when this is true.")
J
Jacek Czaja 已提交
100
        .SetDefault(false);
C
caoying03 已提交
101
    AddComment(R"DOC(
102 103
Softmax Operator.

104
The input of the softmax operator is a tensor of any rank. The output tensor
F
fengjiayi 已提交
105
has the same shape as the input.
C
caoying03 已提交
106

107 108 109 110 111 112
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
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
F
fengjiayi 已提交
113
K-dimensional vector of real values in the range [0, 1] that add up to 1.
114 115 116 117 118
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 已提交
119

F
fengjiayi 已提交
120
For each row $i$ and each column $j$ in the matrix, we have:
F
fengjiayi 已提交
121
    $$Out[i, j] = \frac{\exp(X[i, j])}{\sum_j(exp(X[i, j])}$$
C
caoying03 已提交
122 123

)DOC");
124 125 126
  }
};

C
chengduo 已提交
127 128 129 130 131 132 133 134
class SoftmaxOpInferVarType : public framework::PassInDtypeAndVarTypeToOutput {
 protected:
  std::unordered_map<std::string, std::string> GetInputOutputWithSameType()
      const override {
    return std::unordered_map<std::string, std::string>{{"X", /*->*/ "Out"}};
  }
};

D
dongzhihong 已提交
135
class SoftmaxOpGrad : public framework::OperatorWithKernel {
Y
Yu Yang 已提交
136 137 138
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

139
  void InferShape(framework::InferShapeContext* ctx) const override {
F
fengjiayi 已提交
140 141 142 143 144 145
    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.");
146

147 148
    ctx->SetOutputDim(framework::GradVarName("X"),
                      ctx->GetInputDim(framework::GradVarName("Out")));
D
dongzhihong 已提交
149
  }
150 151 152 153 154

 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override {
    // choose cudnn kernel if the runtime supported.
K
Kexin Zhao 已提交
155
    framework::LibraryType library_{framework::LibraryType::kPlain};
J
Jacek Czaja 已提交
156 157
    std::string data_format = ctx.Attr<std::string>("data_format");
    framework::DataLayout layout_ = framework::StringToDataLayout(data_format);
M
mozga-intel 已提交
158

159
#ifdef PADDLE_WITH_CUDA
K
Kexin Zhao 已提交
160
    if (platform::CanCUDNNBeUsed(ctx)) {
K
Kexin Zhao 已提交
161
      library_ = framework::LibraryType::kCUDNN;
162 163
    }
#endif
J
Jacek Czaja 已提交
164 165 166 167 168 169 170
#ifdef PADDLE_WITH_MKLDNN
    if (library_ == framework::LibraryType::kPlain &&
        platform::CanMKLDNNBeUsed(ctx)) {
      library_ = framework::LibraryType::kMKLDNN;
      layout_ = framework::DataLayout::kMKLDNN;
    }
#endif
Y
Yu Yang 已提交
171 172
    auto input_data_type =
        ctx.Input<Tensor>(framework::GradVarName("Out"))->type();
J
Jacek Czaja 已提交
173 174 175 176 177 178 179
    if (input_data_type == framework::proto::VarType::FP16) {
      PADDLE_ENFORCE(platform::is_gpu_place(ctx.GetPlace()),
                     "float16 can only be used on GPU place");
    }

    return framework::OpKernelType(input_data_type, ctx.GetPlace(), layout_,
                                   library_);
180
  }
D
dongzhihong 已提交
181 182
};

183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200
class SoftmaxOpGradMaker : public framework::SingleGradOpDescMaker {
 public:
  using framework::SingleGradOpDescMaker::SingleGradOpDescMaker;

 protected:
  std::unique_ptr<framework::OpDesc> Apply() const override {
    auto* op = new framework::OpDesc();
    op->SetType("softmax_grad");

    op->SetInput("Out", Output("Out"));
    op->SetInput(framework::GradVarName("Out"), OutputGrad("Out"));

    op->SetAttrMap(Attrs());

    op->SetOutput(framework::GradVarName("X"), InputGrad("X"));
    return std::unique_ptr<framework::OpDesc>(op);
  }
};
D
dzhwinter 已提交
201 202 203 204 205 206 207 208 209 210 211 212 213 214 215

class SoftmaxInplaceInToOut : public framework::InplaceInToOut {
 public:
  using framework::InplaceInToOut::InplaceInToOut;

 protected:
  std::unordered_map<std::string, std::string> Apply(
      const framework::OpDesc& op_desc,
      framework::BlockDesc* block) const override {
    return std::unordered_map<std::string, std::string>{
        {"X", "Out"},
    };
  }
};

216 217 218
}  // namespace operators
}  // namespace paddle

D
dongzhihong 已提交
219
namespace ops = paddle::operators;
D
dongzhihong 已提交
220

Y
Yang Yang 已提交
221
REGISTER_OPERATOR(softmax, ops::SoftmaxOp, ops::SoftmaxOpMaker,
C
chengduo 已提交
222
                  ops::SoftmaxOpInferVarType, ops::SoftmaxOpGradMaker);
223
REGISTER_OPERATOR(softmax_grad, ops::SoftmaxOpGrad);
D
dongzhihong 已提交
224
REGISTER_OP_CPU_KERNEL(
D
dzhwinter 已提交
225 226
    softmax, ops::SoftmaxKernel<paddle::platform::CPUDeviceContext, float>,
    ops::SoftmaxKernel<paddle::platform::CPUDeviceContext, double>);
Q
QI JUN 已提交
227 228
REGISTER_OP_CPU_KERNEL(
    softmax_grad,
D
dzhwinter 已提交
229 230
    ops::SoftmaxGradKernel<paddle::platform::CPUDeviceContext, float>,
    ops::SoftmaxGradKernel<paddle::platform::CPUDeviceContext, double>);