softmax_op.cc 9.2 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 {
Q
Qiao Longfei 已提交
37 38
    PADDLE_ENFORCE(ctx->HasInput("X"),
                   "Input(X) of SoftmaxOp should not be null.");
F
fengjiayi 已提交
39 40
    PADDLE_ENFORCE(ctx->HasOutput("Out"),
                   "Output(Out) of SoftmaxOp should not be null.");
Q
Qiao Longfei 已提交
41

42 43 44
    auto dim_x = ctx->GetInputDim("X");
    auto rank_x = dim_x.size();
    auto axis = ctx->Attrs().Get<int>("axis");
45 46 47 48 49 50
    PADDLE_ENFORCE(axis >= -rank_x && axis < rank_x,
                   "Attr(axis) value should be in range [-R, R-1], "
                   "R is the rank of Input(X).");

    auto use_cudnn = ctx->Attrs().Get<bool>("use_cudnn");
    if (axis != rank_x - 1 && axis != -1) {
D
dengkaipeng 已提交
51
      PADDLE_ENFORCE(!use_cudnn, "CUDNN kernel only support axis as -1.");
52
    }
53

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

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

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

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

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

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

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

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

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

)DOC");
143 144 145
  }
};

C
chengduo 已提交
146 147 148 149 150 151 152 153
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 已提交
154
class SoftmaxOpGrad : public framework::OperatorWithKernel {
Y
Yu Yang 已提交
155 156 157
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

158
  void InferShape(framework::InferShapeContext* ctx) const override {
F
fengjiayi 已提交
159 160 161 162 163 164
    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.");
165

166 167
    ctx->SetOutputDim(framework::GradVarName("X"),
                      ctx->GetInputDim(framework::GradVarName("Out")));
D
dongzhihong 已提交
168
  }
169 170 171 172 173

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

178
#ifdef PADDLE_WITH_CUDA
K
Kexin Zhao 已提交
179
    if (platform::CanCUDNNBeUsed(ctx)) {
K
Kexin Zhao 已提交
180
      library_ = framework::LibraryType::kCUDNN;
181 182
    }
#endif
J
Jacek Czaja 已提交
183 184 185 186 187 188 189
#ifdef PADDLE_WITH_MKLDNN
    if (library_ == framework::LibraryType::kPlain &&
        platform::CanMKLDNNBeUsed(ctx)) {
      library_ = framework::LibraryType::kMKLDNN;
      layout_ = framework::DataLayout::kMKLDNN;
    }
#endif
190 191
    auto input_data_type = OperatorWithKernel::IndicateVarDataType(
        ctx, framework::GradVarName("Out"));
J
Jacek Czaja 已提交
192 193 194 195 196 197 198
    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_);
199
  }
D
dongzhihong 已提交
200 201
};

H
hong 已提交
202 203
template <typename T>
class SoftmaxOpGradMaker : public framework::SingleGradOpMaker<T> {
204
 public:
H
hong 已提交
205
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
206 207

 protected:
208
  void Apply(GradOpPtr<T> op) const override {
209 210
    op->SetType("softmax_grad");

H
hong 已提交
211 212
    op->SetInput("Out", this->Output("Out"));
    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
213

H
hong 已提交
214
    op->SetAttrMap(this->Attrs());
215

H
hong 已提交
216
    op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
217 218
  }
};
D
dzhwinter 已提交
219

220 221
DECLARE_INPLACE_OP_INFERER(SoftmaxInplaceInferer, {"X", "Out"});

222 223 224
// NOTE(zjl): AVX implementation of SoftmaxGrad does not support in-place
DECLARE_CUDA_ONLY_INPLACE_OP_INFERER(SoftmaxGradInplaceInferer,
                                     {"Out", framework::GradVarName("X")});
225

226 227 228
}  // namespace operators
}  // namespace paddle

D
dongzhihong 已提交
229
namespace ops = paddle::operators;
D
dongzhihong 已提交
230

Y
Yang Yang 已提交
231
REGISTER_OPERATOR(softmax, ops::SoftmaxOp, ops::SoftmaxOpMaker,
H
hong 已提交
232 233 234
                  ops::SoftmaxOpInferVarType,
                  ops::SoftmaxOpGradMaker<paddle::framework::OpDesc>,
                  ops::SoftmaxOpGradMaker<paddle::imperative::OpBase>,
235 236 237
                  ops::SoftmaxInplaceInferer);
REGISTER_OPERATOR(softmax_grad, ops::SoftmaxOpGrad,
                  ops::SoftmaxGradInplaceInferer);
D
dongzhihong 已提交
238
REGISTER_OP_CPU_KERNEL(
D
dzhwinter 已提交
239 240
    softmax, ops::SoftmaxKernel<paddle::platform::CPUDeviceContext, float>,
    ops::SoftmaxKernel<paddle::platform::CPUDeviceContext, double>);
Q
QI JUN 已提交
241 242
REGISTER_OP_CPU_KERNEL(
    softmax_grad,
D
dzhwinter 已提交
243 244
    ops::SoftmaxGradKernel<paddle::platform::CPUDeviceContext, float>,
    ops::SoftmaxGradKernel<paddle::platform::CPUDeviceContext, double>);