add autogen code support for margin_cross_entropy (#52130)

paddle/fluid/operators/margin_cross_entropy_op.cc

-/* Copyright (c) 2021 PaddlePaddle Authors. All Rights Reserved.
-
-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
-
-    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. */
-
-#include "paddle/fluid/framework/infershape_utils.h"
-#include "paddle/fluid/framework/op_registry.h"
-#include "paddle/phi/core/infermeta_utils.h"
-#include "paddle/phi/infermeta/backward.h"
-#include "paddle/phi/infermeta/binary.h"
-
-namespace paddle {
-namespace operators {
-
-class MarginCrossEntropyOp : public framework::OperatorWithKernel {
- public:
-  using framework::OperatorWithKernel::OperatorWithKernel;
-
- protected:
-  phi::KernelKey GetExpectedKernelType(
-      const framework::ExecutionContext& ctx) const override {
-    return phi::KernelKey(
-        OperatorWithKernel::IndicateVarDataType(ctx, "Logits"),
-        ctx.device_context().GetPlace());
-  }
-};
-
-class MarginCrossEntropyOpMaker : public framework::OpProtoAndCheckerMaker {
- public:
-  void Make() {
-    AddInput("Logits",
-             "(Tensor, default: Tensor<float>), The input tensor of unscaled "
-             "log probabilities, whose dimension :attr:`axis` should be scaled "
-             "by softmax.");
-    AddInput(
-        "Label",
-        "(Tensor) The input tensor of groud truth label. Label is a "
-        "Tensor<int64> in same shape with Input(Logits) except the shape in "
-        "dimension :attr:`axis` as 1.");
-    AddOutput(
-        "Softmax",
-        "(Tensor, default: Tensor<float>), A tensor in same shape with "
-        "Input(Logits). "
-        "The outputs value of softmax activation by given the input batch, "
-        "which will be used in backward calculation.");
-    AddOutput("Loss",
-              "(Tensor, default: Tensor<float>), A tensor in same shape with "
-              "Input(Logits) "
-              "except the shape in dimension :attr:`axis` as 1. The cross "
-              "entropy loss.");
-    AddAttr<bool>("return_softmax",
-                  "(bool default false) A flag to indicate "
-                  "whether to return softmax.")
-        .SetDefault(false);
-    AddAttr<int>("ring_id", "(int default 0) nccl communication ring id.")
-        .SetDefault(0);
-    AddAttr<int>("rank", "(int default 0) rank id for MarginCrossEntropy.")
-        .SetDefault(0);
-    AddAttr<int>("nranks", "(int default 1) nranks id for MarginCrossEntropy.")
-        .SetDefault(1);
-    AddAttr<float>("margin1", "(float default 1.0) margin1 for MarginLoss.")
-        .SetDefault(1.0);
-    AddAttr<float>("margin2", "(float default 0.5) margin2 for MarginLoss.")
-        .SetDefault(0.5);
-    AddAttr<float>("margin3", "(float default 0.0) margin3 for MarginLoss.")
-        .SetDefault(0.0);
-    AddAttr<float>("scale", "(float default 64.0) scale for MarginLoss.")
-        .SetDefault(64.0);
-    AddComment(R"DOC(
-MarginCrossEntropy Operator
-.. math::
-
-    L=-\frac{1}{N}\sum^N_{i=1}\log\frac{e^{s(cos(m_{1}\theta_{y_i}+m_{2})-m_{3})}}{e^{s(cos(m_{1}\theta_{y_i}+m_{2})-m_{3})}+\sum^n_{j=1,j\neq y_i} e^{scos\theta_{y_i}}}
-
-where the :math: `\theta_{y_i}` is the angle between the feature :math: `x` and
-the representation of class :math: `i`. The details of ArcFace loss
-could be referred to https://arxiv.org/abs/1801.07698.
-
-Note that the Op supports model parallel and single GPU. And Logits.shape[-1] can be different each rank.
-
-)DOC");
-  }
-};
-
-class MarginCrossEntropyOpGrad : public framework::OperatorWithKernel {
- public:
-  using framework::OperatorWithKernel::OperatorWithKernel;
-
- protected:
-  phi::KernelKey GetExpectedKernelType(
-      const framework::ExecutionContext& ctx) const override {
-    return phi::KernelKey(OperatorWithKernel::IndicateVarDataType(
-                              ctx, framework::GradVarName("Loss")),
-                          ctx.device_context().GetPlace());
-  }
-};
-
-template <typename T>
-class MarginCrossEntropyOpGradMaker : public framework::SingleGradOpMaker<T> {
- public:
-  using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
-
- protected:
-  void Apply(GradOpPtr<T> op) const override {
-    op->SetType("margin_cross_entropy_grad");
-
-    op->SetInput("Softmax", this->Output("Softmax"));
-    op->SetInput("Logits", this->Input("Logits"));
-    op->SetInput("Label", this->Input("Label"));
-    op->SetInput(framework::GradVarName("Loss"), this->OutputGrad("Loss"));
-    op->SetAttrMap(this->Attrs());
-    op->SetOutput(framework::GradVarName("Logits"), this->InputGrad("Logits"));
-  }
-};
-
-}  // namespace operators
-}  // namespace paddle
-
-namespace ops = paddle::operators;
-
-DECLARE_INFER_SHAPE_FUNCTOR(margin_cross_entropy,
-                            MarginCrossEntropyInferShapeFunctor,
-                            PD_INFER_META(phi::MarginCrossEntropyInferMeta));
-REGISTER_OPERATOR(
-    margin_cross_entropy,
-    ops::MarginCrossEntropyOp,
-    ops::MarginCrossEntropyOpMaker,
-    ops::MarginCrossEntropyOpGradMaker<paddle::framework::OpDesc>,
-    ops::MarginCrossEntropyOpGradMaker<paddle::imperative::OpBase>,
-    MarginCrossEntropyInferShapeFunctor);
-DECLARE_INFER_SHAPE_FUNCTOR(
-    margin_cross_entropy_grad,
-    MarginCrossEntropyGradInferShapeFunctor,
-    PD_INFER_META(phi::MarginCrossEntropyGradInferMeta));
-REGISTER_OPERATOR(margin_cross_entropy_grad,
-                  ops::MarginCrossEntropyOpGrad,
-                  MarginCrossEntropyGradInferShapeFunctor);

paddle/phi/api/yaml/backward.yaml

@@ -932,6 +932,17 @@
   kernel :
     func : lu_unpack_grad
 
+- backward_op : margin_cross_entropy_grad
+  forward : margin_cross_entropy (Tensor logits, Tensor label, bool return_softmax=false, int ring_id=0, int rank=0, int nranks=1, float margin1=1.0f, float margin2=0.5f, float margin3=0.0f, float scale=64.0f) -> Tensor(softmax), Tensor(loss)
+  args : (Tensor logits, Tensor label, Tensor softmax, Tensor loss_grad, bool return_softmax, int ring_id, int rank, int nranks, float margin1, float margin2, float margin3, float scale)
+  output : Tensor(logits_grad)
+  infer_meta :
+    func : MarginCrossEntropyGradInferMeta
+  kernel :
+    func : margin_cross_entropy_grad
+    data_type : softmax
+  inplace : (softmax -> logits_grad)
+
 - backward_op : masked_select_grad
   forward : masked_select (Tensor x, Tensor mask) -> Tensor(out)
   args : (Tensor x, Tensor mask, Tensor out_grad)

paddle/phi/api/yaml/legacy_backward.yaml

@@ -651,17 +651,6 @@
   kernel :
     func : lu_grad
 
-- backward_op : margin_cross_entropy_grad
-  forward : margin_cross_entropy (Tensor logits, Tensor label, bool return_softmax, int ring_id, int rank, int nranks, float margin1, float margin2, float margin3, float scale) -> Tensor(softmax), Tensor(loss)
-  args : (Tensor logits, Tensor label, Tensor softmax, Tensor loss_grad, bool return_softmax, int ring_id, int rank, int nranks, float margin1, float margin2, float margin3, float scale)
-  output : Tensor(logits_grad)
-  infer_meta :
-    func : MarginCrossEntropyGradInferMeta
-  kernel :
-    func : margin_cross_entropy_grad
-    data_type : softmax
-  inplace : (softmax -> logits_grad)
-
 - backward_op : matmul_double_grad
   forward : matmul_grad (Tensor x, Tensor y, Tensor grad_out, bool transpose_x=false, bool transpose_y=false) -> Tensor(grad_x), Tensor(grad_y)
   args : (Tensor x, Tensor y, Tensor grad_out, Tensor grad_x_grad, Tensor grad_y_grad, bool transpose_x=false, bool transpose_y=false)

paddle/phi/api/yaml/legacy_ops.yaml

@@ -972,16 +972,6 @@
     func : lu
   backward : lu_grad
 
-- op : margin_cross_entropy
-  args : (Tensor logits, Tensor label, bool return_softmax, int ring_id, int rank, int nranks, float margin1, float margin2, float margin3, float scale)
-  output : Tensor(softmax), Tensor(loss)
-  infer_meta :
-    func : MarginCrossEntropyInferMeta
-  kernel :
-    func : margin_cross_entropy
-    data_type : logits
-  backward : margin_cross_entropy_grad
-
 - op : matmul
   args : (Tensor x, Tensor y, bool transpose_x = false, bool transpose_y = false)
   output : Tensor

paddle/phi/api/yaml/op_compat.yaml

@@ -1137,6 +1137,13 @@
   outputs :
     {pmat : Pmat, l : L, u : U}
 
+- op : margin_cross_entropy
+  backward : margin_cross_entropy_grad
+  inputs:
+    {logits : Logits, label : Label}
+  outputs:
+    {softmax : Softmax, loss : Loss}
+
 - op : masked_select
   inputs :
     {x : X, mask : Mask}

paddle/phi/api/yaml/ops.yaml

@@ -933,6 +933,16 @@
     data_type : x
   backward : lu_unpack_grad
 
+- op : margin_cross_entropy
+  args : (Tensor logits, Tensor label, bool return_softmax = false, int ring_id = 0, int rank = 0, int nranks = 1, float margin1 = 1.0f, float margin2 = 0.5f, float margin3 = 0.0f, float scale = 64.0f)
+  output : Tensor(softmax), Tensor(loss)
+  infer_meta :
+    func : MarginCrossEntropyInferMeta
+  kernel :
+    func : margin_cross_entropy
+    data_type : logits
+  backward : margin_cross_entropy_grad
+
 - op : masked_select
   args : (Tensor x, Tensor mask)
   output : Tensor (out)

paddle/phi/ops/compat/margin_cross_entropy_sig.cc

-// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
-//
-// 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
-//
-//     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.
-
-#include "paddle/phi/core/compat/op_utils.h"
-
-namespace phi {
-
-KernelSignature MarginCrossEntropyOpArgumentMapping(
-    const ArgumentMappingContext& ctx) {
-  return KernelSignature("margin_cross_entropy",
-                         {"Logits", "Label"},
-                         {"return_softmax",
-                          "ring_id",
-                          "rank",
-                          "nranks",
-                          "margin1",
-                          "margin2",
-                          "margin3",
-                          "scale"},
-                         {"Softmax", "Loss"});
-}
-
-KernelSignature MarginCrossEntropyGradOpArgumentMapping(
-    const ArgumentMappingContext& ctx) {
-  return KernelSignature("margin_cross_entropy_grad",
-                         {"Logits", "Label", "Softmax", "Loss@GRAD"},
-                         {"return_softmax",
-                          "ring_id",
-                          "rank",
-                          "nranks",
-                          "margin1",
-                          "margin2",
-                          "margin3",
-                          "scale"},
-                         {"Logits@GRAD"});
-}
-
-}  // namespace phi
-
-PD_REGISTER_ARG_MAPPING_FN(margin_cross_entropy,
-                           phi::MarginCrossEntropyOpArgumentMapping);
-PD_REGISTER_ARG_MAPPING_FN(margin_cross_entropy_grad,
-                           phi::MarginCrossEntropyGradOpArgumentMapping);