softmax_op.cc 9.8 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

L
liuwei1031 已提交
17
#include <memory>
18
#include <string>
L
liuwei1031 已提交
19
#include <unordered_map>
20

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

25 26 27
#ifdef PADDLE_WITH_MKLDNN
#include "paddle/fluid/platform/mkldnn_helper.h"
#endif
28

29 30 31
namespace paddle {
namespace operators {

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

36
  void InferShape(framework::InferShapeContext* ctx) const override {
37 38 39 40 41 42
    PADDLE_ENFORCE_EQ(
        ctx->HasInput("X"), true,
        platform::errors::NotFound("Input(X) of SoftmaxOp is not found."));
    PADDLE_ENFORCE_EQ(
        ctx->HasOutput("Out"), true,
        platform::errors::NotFound("Output(Out) of SoftmaxOp is not found."));
Q
Qiao Longfei 已提交
43

44 45 46
    auto dim_x = ctx->GetInputDim("X");
    auto rank_x = dim_x.size();
    auto axis = ctx->Attrs().Get<int>("axis");
47 48 49 50 51 52 53 54
    PADDLE_ENFORCE_GE(axis, -rank_x,
                      platform::errors::InvalidArgument(
                          "Attr(axis) value should be in range [-R, R-1], "
                          "R is the rank of Input(X)."));
    PADDLE_ENFORCE_LT(axis, rank_x,
                      platform::errors::InvalidArgument(
                          "Attr(axis) value should be in range [-R, R-1], "
                          "R is the rank of Input(X)."));
55

F
fengjiayi 已提交
56
    ctx->SetOutputDim("Out", ctx->GetInputDim("X"));
Q
Qiao Longfei 已提交
57
    ctx->ShareLoD("X", /*->*/ "Out");
58
  }
59 60 61 62 63

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

68
#ifdef PADDLE_WITH_CUDA
K
Kexin Zhao 已提交
69
    if (platform::CanCUDNNBeUsed(ctx)) {
K
Kexin Zhao 已提交
70
      library_ = framework::LibraryType::kCUDNN;
71 72
    }
#endif
73 74 75 76
#ifdef PADDLE_WITH_MKLDNN
    if (library_ == framework::LibraryType::kPlain &&
        platform::CanMKLDNNBeUsed(ctx)) {
      library_ = framework::LibraryType::kMKLDNN;
M
mozga-intel 已提交
77
      layout_ = framework::DataLayout::kMKLDNN;
78 79
    }
#endif
K
Kexin Zhao 已提交
80

81
    auto input_data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X");
K
Kexin Zhao 已提交
82
    if (input_data_type == framework::proto::VarType::FP16) {
83 84 85
      PADDLE_ENFORCE_EQ(platform::is_gpu_place(ctx.GetPlace()), true,
                        platform::errors::InvalidArgument(
                            "float16 can only be used on GPU place"));
K
Kexin Zhao 已提交
86 87
    }

M
mozga-intel 已提交
88
    return framework::OpKernelType(input_data_type, ctx.GetPlace(), layout_,
K
Kexin Zhao 已提交
89
                                   library_);
90
  }
91
};
92

D
dongzhihong 已提交
93
class SoftmaxOpMaker : public framework::OpProtoAndCheckerMaker {
94
 public:
Y
Yu Yang 已提交
95
  void Make() override {
96
    AddInput("X",
F
fengjiayi 已提交
97
             "The input tensor of softmax, "
D
dengkaipeng 已提交
98
             "whose dimension :attr:`axis` is the input_feature_dimensions.");
99
    AddOutput("Out", "The normalized values with the same shape as X.");
100
    AddAttr<int>("axis",
D
dengkaipeng 已提交
101
                 "The dimension index of Input(x) to perform softmax,"
102 103
                 "default -1 for last dimension")
        .SetDefault(-1);
104 105 106 107 108 109 110 111 112 113 114
    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");
115 116 117
    AddAttr<bool>("use_mkldnn",
                  "(bool, default false) Only used in mkldnn kernel")
        .SetDefault(false);
118 119 120 121 122
    AddAttr<std::string>(
        "mkldnn_data_type",
        "(string, default \"float32\"). Data type of mkldnn kernel")
        .SetDefault("float32")
        .InEnum({"float32", "bfloat16"});
J
Jacek Czaja 已提交
123
    AddAttr<bool>("is_test",
124 125
                  "(bool, default false) Set to true for inference only, false "
                  "for training. Some layers may run faster when this is true.")
J
Jacek Czaja 已提交
126
        .SetDefault(false);
C
caoying03 已提交
127
    AddComment(R"DOC(
128 129
Softmax Operator.

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

D
dengkaipeng 已提交
133
The dimension :attr:`axis` of the input tensor will be permuted to the last.
D
dengkaipeng 已提交
134
Then the input tensor will be logically flattened to a 2-D matrix. The matrix's
D
dengkaipeng 已提交
135
second dimension(row length) is as same as the dimension :attr:`axis` of the input
136 137 138
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
D
dengkaipeng 已提交
139
of the input tensor's dimension :attr:`axis`) vector of arbitrary real values to a
F
fengjiayi 已提交
140
K-dimensional vector of real values in the range [0, 1] that add up to 1.
141 142 143 144 145
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 已提交
146

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

)DOC");
151 152 153
  }
};

C
chengduo 已提交
154 155
class SoftmaxOpInferVarType : public framework::PassInDtypeAndVarTypeToOutput {
 protected:
156
  std::unordered_map<std::string, std::string>& GetInputOutputWithSameType()
C
chengduo 已提交
157
      const override {
158 159
    static std::unordered_map<std::string, std::string> m{{"X", /*->*/ "Out"}};
    return m;
C
chengduo 已提交
160 161 162
  }
};

D
dongzhihong 已提交
163
class SoftmaxOpGrad : public framework::OperatorWithKernel {
Y
Yu Yang 已提交
164 165 166
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

167
  void InferShape(framework::InferShapeContext* ctx) const override {
168 169 170 171 172 173 174 175 176 177 178
    PADDLE_ENFORCE_EQ(
        ctx->HasInput("Out"), true,
        platform::errors::InvalidArgument("Input(Out) is not found."));
    PADDLE_ENFORCE_EQ(
        ctx->HasInput(framework::GradVarName("Out")), true,
        platform::errors::InvalidArgument("Input(Out@GRAD) is not found."));
    PADDLE_ENFORCE_EQ(
        ctx->GetInputDim("Out"),
        ctx->GetInputDim(framework::GradVarName("Out")),
        platform::errors::InvalidArgument("Input(Out) and its gradients "
                                          "should have a same shape."));
179

180 181
    ctx->SetOutputDim(framework::GradVarName("X"),
                      ctx->GetInputDim(framework::GradVarName("Out")));
D
dongzhihong 已提交
182
  }
183 184 185 186 187

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

192
#ifdef PADDLE_WITH_CUDA
K
Kexin Zhao 已提交
193
    if (platform::CanCUDNNBeUsed(ctx)) {
K
Kexin Zhao 已提交
194
      library_ = framework::LibraryType::kCUDNN;
195 196
    }
#endif
J
Jacek Czaja 已提交
197 198 199 200 201 202 203
#ifdef PADDLE_WITH_MKLDNN
    if (library_ == framework::LibraryType::kPlain &&
        platform::CanMKLDNNBeUsed(ctx)) {
      library_ = framework::LibraryType::kMKLDNN;
      layout_ = framework::DataLayout::kMKLDNN;
    }
#endif
204 205
    auto input_data_type = OperatorWithKernel::IndicateVarDataType(
        ctx, framework::GradVarName("Out"));
J
Jacek Czaja 已提交
206
    if (input_data_type == framework::proto::VarType::FP16) {
207 208 209
      PADDLE_ENFORCE_EQ(platform::is_gpu_place(ctx.GetPlace()), true,
                        platform::errors::InvalidArgument(
                            "float16 can only be used on GPU place"));
J
Jacek Czaja 已提交
210 211 212 213
    }

    return framework::OpKernelType(input_data_type, ctx.GetPlace(), layout_,
                                   library_);
214
  }
D
dongzhihong 已提交
215 216
};

H
hong 已提交
217 218
template <typename T>
class SoftmaxOpGradMaker : public framework::SingleGradOpMaker<T> {
219
 public:
H
hong 已提交
220
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
221 222

 protected:
223
  void Apply(GradOpPtr<T> op) const override {
224 225
    op->SetType("softmax_grad");

H
hong 已提交
226 227
    op->SetInput("Out", this->Output("Out"));
    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
228

H
hong 已提交
229
    op->SetAttrMap(this->Attrs());
230

H
hong 已提交
231
    op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
232 233
  }
};
D
dzhwinter 已提交
234

235 236
DECLARE_INPLACE_OP_INFERER(SoftmaxInplaceInferer, {"X", "Out"});

237 238 239
// NOTE(zjl): AVX implementation of SoftmaxGrad does not support in-place
DECLARE_CUDA_ONLY_INPLACE_OP_INFERER(SoftmaxGradInplaceInferer,
                                     {"Out", framework::GradVarName("X")});
240

241 242 243
}  // namespace operators
}  // namespace paddle

D
dongzhihong 已提交
244
namespace ops = paddle::operators;
D
dongzhihong 已提交
245

Y
Yang Yang 已提交
246
REGISTER_OPERATOR(softmax, ops::SoftmaxOp, ops::SoftmaxOpMaker,
H
hong 已提交
247 248 249
                  ops::SoftmaxOpInferVarType,
                  ops::SoftmaxOpGradMaker<paddle::framework::OpDesc>,
                  ops::SoftmaxOpGradMaker<paddle::imperative::OpBase>,
250 251 252
                  ops::SoftmaxInplaceInferer);
REGISTER_OPERATOR(softmax_grad, ops::SoftmaxOpGrad,
                  ops::SoftmaxGradInplaceInferer);
D
dongzhihong 已提交
253
REGISTER_OP_CPU_KERNEL(
D
dzhwinter 已提交
254 255
    softmax, ops::SoftmaxKernel<paddle::platform::CPUDeviceContext, float>,
    ops::SoftmaxKernel<paddle::platform::CPUDeviceContext, double>);
Q
QI JUN 已提交
256 257
REGISTER_OP_CPU_KERNEL(
    softmax_grad,
D
dzhwinter 已提交
258 259
    ops::SoftmaxGradKernel<paddle::platform::CPUDeviceContext, float>,
    ops::SoftmaxGradKernel<paddle::platform::CPUDeviceContext, double>);