未验证 提交 d0514a93 编写于 作者: L LoneRanger 提交者: GitHub

[static op generation] group_norm (#53489)

* fix the static op generation for group_norm

* fix bug of mismatch

* fix bug of AssertionError

* fix setting of composite
上级 3ac1ccf9
/* Copyright (c) 2018 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 <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "paddle/fluid/prim/api/composite_backward/composite_backward_api.h"
#include "paddle/fluid/prim/utils/static/composite_grad_desc_maker.h"
#include "paddle/fluid/prim/utils/static/desc_tensor.h"
#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/ternary.h"
namespace paddle {
namespace operators {
using DataLayout = phi::DataLayout;
class GroupNormOp : public framework::OperatorWithKernel {
public:
using framework::OperatorWithKernel::OperatorWithKernel;
};
class GroupNormOpMaker : public framework::OpProtoAndCheckerMaker {
public:
void Make() override {
AddInput("X", "The input tensor.");
AddInput("Scale",
"Scale is a 1-dimensional tensor of size C"
"that is applied to the output.")
.AsDispensable();
AddInput("Bias",
"Bias is a 1-dimensional tensor of size C "
"that is applied to the output")
.AsDispensable();
AddOutput("Y", "Result after normalization.");
AddOutput("Mean", "Mean of each group.").AsIntermediate();
AddOutput("Variance", "Variance of each group.").AsIntermediate();
AddAttr<float>("epsilon",
"Constant for numerical stability [default 1e-5].")
.SetDefault(1e-5)
.AddCustomChecker([](const float &epsilon) {
PADDLE_ENFORCE_EQ(epsilon >= 0.0f && epsilon <= 1.0f,
true,
platform::errors::InvalidArgument(
"'epsilon' in Op(GroupNorm) should be between"
"0.0 and 1.0f, But received [%s].",
epsilon));
});
AddAttr<int>("groups", "The number of groups that divided from channels.")
.AddCustomChecker([](const int &groups) {
PADDLE_ENFORCE_GT(
groups,
0,
platform::errors::InvalidArgument(
"'groups' in Op(GroupNorm) should be greater than zero,"
"But received [%s].",
groups));
});
AddAttr<std::string>("data_layout",
"An optional string from: \"NHWC\", \"NCHW\". ")
.SetDefault("NCHW");
AddComment(R"DOC(
Group Normalization
Refer to `Group Normalization <https://arxiv.org/abs/1803.08494>`_
)DOC");
}
};
class GroupNormGradOp : public framework::OperatorWithKernel {
public:
using framework::OperatorWithKernel::OperatorWithKernel;
void InferShape(framework::InferShapeContext *ctx) const override {
// check input
OP_INOUT_CHECK(ctx->HasInput("X"), "Input", "X", "GroupNormGrad");
OP_INOUT_CHECK(ctx->HasInput("Y"), "Input", "Y", "GroupNormGrad");
OP_INOUT_CHECK(
ctx->HasInput("Variance"), "Input", "Variance", "GroupNormGrad");
OP_INOUT_CHECK(ctx->HasInput("Mean"), "Input", "Mean", "GroupNormGrad");
OP_INOUT_CHECK(ctx->HasInput(framework::GradVarName("Y")),
"Input",
framework::GradVarName("Y"),
"GroupNormGrad");
// check output
if (ctx->HasOutput(framework::GradVarName("X"))) {
ctx->SetOutputDim(framework::GradVarName("X"), ctx->GetInputDim("Y"));
}
if (ctx->HasOutput(framework::GradVarName("Scale"))) {
ctx->SetOutputDim(framework::GradVarName("Scale"),
ctx->GetInputDim("Scale"));
}
if (ctx->HasOutput(framework::GradVarName("Bias"))) {
ctx->SetOutputDim(framework::GradVarName("Bias"),
ctx->GetInputDim("Bias"));
}
}
protected:
phi::KernelKey GetExpectedKernelType(
const framework::ExecutionContext &ctx) const override {
const auto *var = ctx.InputVar(framework::GradVarName("Y"));
PADDLE_ENFORCE_NOT_NULL(
var,
platform::errors::InvalidArgument(
"Input(Y@GRAD) of GroupNormGradOp should not be null"));
const phi::DenseTensor *t = nullptr;
if (var->IsType<phi::DenseTensor>()) {
t = &var->Get<phi::DenseTensor>();
} else if (var->IsType<phi::DenseTensor>()) {
t = &var->Get<phi::DenseTensor>();
}
PADDLE_ENFORCE_NOT_NULL(t,
platform::errors::InvalidArgument(
"Input(Y@GRAD) phi::DenseTensor of "
"GroupNormGradOp should not be null"));
return phi::KernelKey(framework::TransToProtoVarType(t->dtype()),
ctx.GetPlace());
}
};
template <typename T>
class GroupNormGradMaker : public framework::SingleGradOpMaker<T> {
public:
using framework::SingleGradOpMaker<T>::SingleGradOpMaker;
void Apply(GradOpPtr<T> op) const override {
op->SetType("group_norm_grad");
op->SetInput("X", this->Input("X"));
op->SetInput("Scale", this->Input("Scale"));
op->SetInput("Bias", this->Input("Bias"));
op->SetInput(framework::GradVarName("Y"), this->OutputGrad("Y"));
op->SetInput("Y", this->Output("Y"));
op->SetInput("Mean", this->Output("Mean"));
op->SetInput("Variance", this->Output("Variance"));
op->SetOutput(framework::GradVarName("X"), this->InputGrad("X"));
op->SetOutput(framework::GradVarName("Bias"), this->InputGrad("Bias"));
op->SetOutput(framework::GradVarName("Scale"), this->InputGrad("Scale"));
op->SetAttrMap(this->Attrs());
}
};
class GroupNormCompositeGradOpMaker : public prim::CompositeGradOpMakerBase {
using prim::CompositeGradOpMakerBase::CompositeGradOpMakerBase;
public:
void Apply() override {
// inputs and outputs of group_norm
paddle::Tensor x = this->GetSingleForwardInput("X");
paddle::optional<paddle::Tensor> scale =
this->GetOptionalSingleForwardInput("Scale");
paddle::optional<paddle::Tensor> bias =
this->GetOptionalSingleForwardInput("Bias");
paddle::Tensor y = this->GetSingleForwardOutput("Y");
paddle::Tensor mean = this->GetSingleForwardOutput("Mean");
paddle::Tensor variance = this->GetSingleForwardOutput("Variance");
paddle::Tensor y_grad = this->GetSingleOutputGrad("Y");
paddle::Tensor x_grad = this->GetSingleInputGrad("X");
paddle::Tensor scale_grad = this->GetSingleInputGrad("Scale");
paddle::Tensor bias_grad = this->GetSingleInputGrad("Bias");
auto dx_ptr = this->GetOutputPtr(&x_grad);
std::string dx_name = this->GetOutputName(x_grad);
auto dscale_ptr = this->GetOutputPtr(&scale_grad);
std::string dscale_name = this->GetOutputName(scale_grad);
auto dbias_ptr = this->GetOutputPtr(&bias_grad);
std::string dbias_name = this->GetOutputName(bias_grad);
// attrs of group_norm
auto groups = this->Attr<int>("groups");
auto epsilon = this->Attr<float>("epsilon");
auto data_layout = this->Attr<std::string>("data_layout");
VLOG(3) << "Runing group_norm composite func";
prim::group_norm_grad<prim::DescTensor>(x,
scale,
bias,
y,
mean,
variance,
y_grad,
epsilon,
groups,
data_layout,
dx_ptr,
dscale_ptr,
dbias_ptr);
this->RecoverOutputName(x_grad, dx_name);
this->RecoverOutputName(scale_grad, dscale_name);
this->RecoverOutputName(bias_grad, dbias_name);
}
};
DECLARE_INPLACE_OP_INFERER(GroupNormGradInplaceInferer,
{framework::GradVarName("Y"),
framework::GradVarName("X")});
class GroupNormOpInferVarType
: public framework::PassInDtypeAndVarTypeToOutput {
protected:
std::unordered_map<std::string, std::string> &GetInputOutputWithSameType()
const override {
static std::unordered_map<std::string, std::string> m{{"X", /*->*/ "Y"}};
return m;
}
};
} // namespace operators
} // namespace paddle
DECLARE_INFER_SHAPE_FUNCTOR(group_norm,
GroupNormInferShapeFunctor,
PD_INFER_META(phi::GroupNormInferMeta));
namespace ops = paddle::operators;
REGISTER_OPERATOR(group_norm,
ops::GroupNormOp,
ops::GroupNormOpMaker,
ops::GroupNormOpInferVarType,
ops::GroupNormGradMaker<paddle::framework::OpDesc>,
ops::GroupNormGradMaker<paddle::imperative::OpBase>,
ops::GroupNormCompositeGradOpMaker,
GroupNormInferShapeFunctor);
REGISTER_OPERATOR(group_norm_grad,
ops::GroupNormGradOp,
ops::GroupNormGradInplaceInferer);
...@@ -749,6 +749,20 @@ ...@@ -749,6 +749,20 @@
func : grid_sample_grad func : grid_sample_grad
data_type : x data_type : x
- backward_op : group_norm_grad
forward : group_norm (Tensor x, Tensor scale, Tensor bias, float epsilon = 1e-5, int groups = -1, str data_layout = "NCHW") -> Tensor(y), Tensor(mean), Tensor(variance)
args : (Tensor x, Tensor scale, Tensor bias, Tensor y, Tensor mean, Tensor variance, Tensor y_grad, float epsilon, int groups, str data_layout)
output : Tensor(x_grad), Tensor(scale_grad), Tensor(bias_grad)
infer_meta :
func : GeneralTernaryGradInferMeta
param : [y, scale, bias]
kernel :
func : group_norm_grad
data_type : y_grad
composite : group_norm_grad(x, scale, bias, y, mean, variance, y_grad, epsilon, groups, data_layout, x_grad, scale_grad, bias_grad)
optional: scale, bias
inplace : (y_grad -> x_grad)
- backward_op : gumbel_softmax_grad - backward_op : gumbel_softmax_grad
forward : gumbel_softmax (Tensor x, float temperature, bool hard, int axis) -> Tensor(out) forward : gumbel_softmax (Tensor x, float temperature, bool hard, int axis) -> Tensor(out)
args : (Tensor out, Tensor out_grad, int axis) args : (Tensor out, Tensor out_grad, int axis)
......
...@@ -408,20 +408,6 @@ ...@@ -408,20 +408,6 @@
composite : gather_grad(x, index, out_grad, axis, x_grad) composite : gather_grad(x, index, out_grad, axis, x_grad)
no_need_buffer : x no_need_buffer : x
- backward_op : group_norm_grad
forward : group_norm (Tensor x, Tensor scale, Tensor bias, float epsilon, int groups, str data_layout) -> Tensor(y), Tensor(mean), Tensor(variance)
args : (Tensor x, Tensor scale, Tensor bias, Tensor y, Tensor mean, Tensor variance, Tensor y_grad, float epsilon, int groups, str data_layout)
output : Tensor(x_grad), Tensor(scale_grad), Tensor(bias_grad)
infer_meta :
func : GeneralTernaryGradInferMeta
param : [y, scale, bias]
kernel :
func : group_norm_grad
data_type : y_grad
composite : group_norm_grad(x, scale, bias, y, mean, variance, y_grad, epsilon, groups, data_layout)
optional: scale, bias
inplace : (y_grad -> x_grad)
- backward_op : hardswish_grad - backward_op : hardswish_grad
forward : hardswish (Tensor x) -> Tensor(out) forward : hardswish (Tensor x) -> Tensor(out)
args : (Tensor x, Tensor out_grad) args : (Tensor x, Tensor out_grad)
......
...@@ -538,17 +538,6 @@ ...@@ -538,17 +538,6 @@
kernel : kernel :
func : greater_than func : greater_than
- op : group_norm
args : (Tensor x, Tensor scale, Tensor bias, float epsilon, int groups, str data_layout)
output : Tensor(y), Tensor(mean), Tensor(variance)
infer_meta :
func : GroupNormInferMeta
kernel :
func : group_norm
optional : scale, bias
intermediate : mean, variance
backward : group_norm_grad
- op : hardswish - op : hardswish
args : (Tensor x) args : (Tensor x)
output : Tensor(out) output : Tensor(out)
......
...@@ -873,6 +873,17 @@ ...@@ -873,6 +873,17 @@
data_type : x data_type : x
backward : grid_sample_grad backward : grid_sample_grad
- op : group_norm
args : (Tensor x, Tensor scale, Tensor bias, float epsilon = 1e-5, int groups = -1, str data_layout = "NCHW")
output : Tensor(y), Tensor(mean), Tensor(variance)
infer_meta :
func : GroupNormInferMeta
kernel :
func : group_norm
optional : scale, bias
intermediate : mean, variance
backward : group_norm_grad
- op : gumbel_softmax - op : gumbel_softmax
args : (Tensor x, float temperature = 1.0, bool hard = false, int axis = -1) args : (Tensor x, float temperature = 1.0, bool hard = false, int axis = -1)
output : Tensor output : Tensor
......
// 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 GroupNormOpArgumentMapping(
const ArgumentMappingContext& ctx UNUSED) {
return KernelSignature("group_norm",
{"X", "Scale", "Bias"},
{"epsilon", "groups", "data_layout"},
{"Y", "Mean", "Variance"});
}
KernelSignature GroupNormGradOpArgumentMapping(
const ArgumentMappingContext& ctx UNUSED) {
return KernelSignature(
"group_norm_grad",
{"X", "Scale", "Bias", "Y", "Mean", "Variance", "Y@GRAD"},
{"epsilon", "groups", "data_layout"},
{"X@GRAD", "Scale@GRAD", "Bias@GRAD"});
}
} // namespace phi
PD_REGISTER_ARG_MAPPING_FN(group_norm, phi::GroupNormOpArgumentMapping);
PD_REGISTER_ARG_MAPPING_FN(group_norm_grad,
phi::GroupNormGradOpArgumentMapping);
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册