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

L
liuwei1031 已提交
15
#include <memory>
16
#include <string>
L
liuwei1031 已提交
17
#include <unordered_map>
18

19
#include "paddle/fluid/framework/infershape_utils.h"
20
#include "paddle/fluid/framework/op_registry.h"
21
#include "paddle/fluid/platform/device/gpu/gpu_dnn.h"
22

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

27 28 29 30
#include "paddle/phi/core/infermeta_utils.h"
#include "paddle/phi/infermeta/backward.h"
#include "paddle/phi/infermeta/unary.h"

31 32 33
namespace paddle {
namespace operators {

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

38 39 40 41
 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override {
    // choose cudnn kernel if the runtime supported.
M
mozga-intel 已提交
42
    std::string data_format = ctx.Attr<std::string>("data_format");
43
    phi::DataLayout layout_ = phi::StringToDataLayout(data_format);
44
    auto input_data_type = OperatorWithKernel::IndicateVarDataType(ctx, "X");
K
Kexin Zhao 已提交
45
    if (input_data_type == framework::proto::VarType::FP16) {
46 47 48 49
      PADDLE_ENFORCE_EQ(
          platform::is_gpu_place(ctx.GetPlace()) ||
              platform::is_npu_place(ctx.GetPlace()) ||
              platform::is_xpu_place(ctx.GetPlace()) ||
50 51
              platform::is_mlu_place(ctx.GetPlace()) ||
              platform::is_custom_place(ctx.GetPlace()),
52 53
          true,
          platform::errors::InvalidArgument(
54
              "float16 can only be used on GPU/NPU/XPU/MLU and custom place"));
K
Kexin Zhao 已提交
55
    }
H
HongyuJia 已提交
56 57 58 59 60 61 62 63 64
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
    if (platform::CanCUDNNBeUsed(ctx)) {
      return framework::OpKernelType(input_data_type,
                                     ctx.GetPlace(),
                                     layout_,
                                     framework::LibraryType::kCUDNN);
    }
#endif
    return framework::OpKernelType(input_data_type, ctx.GetPlace(), layout_);
65
  }
66
};
67

D
dongzhihong 已提交
68
class SoftmaxOpMaker : public framework::OpProtoAndCheckerMaker {
69
 public:
Y
Yu Yang 已提交
70
  void Make() override {
71
    AddInput("X",
F
fengjiayi 已提交
72
             "The input tensor of softmax, "
D
dengkaipeng 已提交
73
             "whose dimension :attr:`axis` is the input_feature_dimensions.");
74
    AddOutput("Out", "The normalized values with the same shape as X.");
75
    AddAttr<int>("axis",
D
dengkaipeng 已提交
76
                 "The dimension index of Input(x) to perform softmax,"
77 78
                 "default -1 for last dimension")
        .SetDefault(-1);
79 80 81 82 83 84 85
    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");
86 87 88 89 90
    AddAttr<bool>(
        "use_cudnn",
        "(bool, default false) Only used in cudnn kernel, need install cudnn")
        .SetDefault(false)
        .AsExtra();
C
caoying03 已提交
91
    AddComment(R"DOC(
92 93
Softmax Operator.

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

D
dengkaipeng 已提交
97
The dimension :attr:`axis` of the input tensor will be permuted to the last.
D
dengkaipeng 已提交
98
Then the input tensor will be logically flattened to a 2-D matrix. The matrix's
D
dengkaipeng 已提交
99
second dimension(row length) is as same as the dimension :attr:`axis` of the input
100 101 102
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 已提交
103
of the input tensor's dimension :attr:`axis`) vector of arbitrary real values to a
F
fengjiayi 已提交
104
K-dimensional vector of real values in the range [0, 1] that add up to 1.
105 106 107 108 109
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 已提交
110

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

)DOC");
115 116 117
  }
};

C
chengduo 已提交
118 119
class SoftmaxOpInferVarType : public framework::PassInDtypeAndVarTypeToOutput {
 protected:
120
  std::unordered_map<std::string, std::string>& GetInputOutputWithSameType()
C
chengduo 已提交
121
      const override {
122 123
    static std::unordered_map<std::string, std::string> m{{"X", /*->*/ "Out"}};
    return m;
C
chengduo 已提交
124 125 126
  }
};

D
dongzhihong 已提交
127
class SoftmaxOpGrad : public framework::OperatorWithKernel {
Y
Yu Yang 已提交
128 129 130
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

131 132 133 134
 protected:
  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override {
    // choose cudnn kernel if the runtime supported.
J
Jacek Czaja 已提交
135
    std::string data_format = ctx.Attr<std::string>("data_format");
136
    phi::DataLayout layout_ = phi::StringToDataLayout(data_format);
137 138
    auto input_data_type = OperatorWithKernel::IndicateVarDataType(
        ctx, framework::GradVarName("Out"));
J
Jacek Czaja 已提交
139
    if (input_data_type == framework::proto::VarType::FP16) {
140
      if (!(platform::is_gpu_place(ctx.GetPlace()) ||
141
            platform::is_npu_place(ctx.GetPlace()) ||
142
            platform::is_xpu_place(ctx.GetPlace()) ||
143 144
            platform::is_mlu_place(ctx.GetPlace()) ||
            platform::is_custom_place(ctx.GetPlace())))
145
        PADDLE_THROW(platform::errors::InvalidArgument(
146
            "float16 can only be used on GPU/NPU/XPU/MLU and custom place"));
J
Jacek Czaja 已提交
147
    }
H
HongyuJia 已提交
148 149 150 151 152 153 154 155 156
#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
    if (platform::CanCUDNNBeUsed(ctx)) {
      return framework::OpKernelType(input_data_type,
                                     ctx.GetPlace(),
                                     layout_,
                                     framework::LibraryType::kCUDNN);
    }
#endif
    return framework::OpKernelType(input_data_type, ctx.GetPlace(), layout_);
157
  }
D
dongzhihong 已提交
158 159
};

H
hong 已提交
160 161
template <typename T>
class SoftmaxOpGradMaker : public framework::SingleGradOpMaker<T> {
162
 public:
H
hong 已提交
163
  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
164 165

 protected:
166
  void Apply(GradOpPtr<T> op) const override {
167 168
    op->SetType("softmax_grad");

H
hong 已提交
169 170
    op->SetInput("Out", this->Output("Out"));
    op->SetInput(framework::GradVarName("Out"), this->OutputGrad("Out"));
171

H
hong 已提交
172
    op->SetAttrMap(this->Attrs());
173

H
hong 已提交
174
    op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
175 176
  }
};
D
dzhwinter 已提交
177

178 179
DECLARE_INPLACE_OP_INFERER(SoftmaxInplaceInferer, {"X", "Out"});

180 181 182
}  // namespace operators
}  // namespace paddle

D
dongzhihong 已提交
183
namespace ops = paddle::operators;
D
dongzhihong 已提交
184

185 186
DECLARE_INFER_SHAPE_FUNCTOR(softmax,
                            SoftmaxInferShapeFunctor,
187
                            PD_INFER_META(phi::SoftmaxInferMeta));
188 189 190
REGISTER_OPERATOR(softmax,
                  ops::SoftmaxOp,
                  ops::SoftmaxOpMaker,
H
hong 已提交
191 192 193
                  ops::SoftmaxOpInferVarType,
                  ops::SoftmaxOpGradMaker<paddle::framework::OpDesc>,
                  ops::SoftmaxOpGradMaker<paddle::imperative::OpBase>,
194 195 196 197
                  ops::SoftmaxInplaceInferer,
                  SoftmaxInferShapeFunctor);
DECLARE_INFER_SHAPE_FUNCTOR(softmax_grad,
                            SoftmaxGradInferShapeFunctor,
198
                            PD_INFER_META(phi::GeneralUnaryGradInferMeta));
199 200
REGISTER_OPERATOR(softmax_grad,
                  ops::SoftmaxOpGrad,
C
Chen Weihang 已提交
201
                  SoftmaxGradInferShapeFunctor);