Add of GlobalMaxPool Gradient (#23502)

### Description
Added gradient computation support for the GlobalMaxPool node.



### Motivation and Context
Improve the training capabilities of ONNX Runtime.
This commit is contained in:
Corentin Maravat 2025-01-28 18:00:01 +01:00 committed by GitHub
parent ded8730d6e
commit d2c5e2474c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 43 additions and 0 deletions

View file

@ -2239,5 +2239,23 @@ IMPLEMENT_GRADIENT_BUILDER(GetAtanGradient) {
return result;
}
IMPLEMENT_GRADIENT_BUILDER(GetGlobalMaxPoolGradient) {
// For GlobalMaxPool's gradient, a binary mask flags max elements.
// We multiply that mask by the incoming gradient, passing gradients only to maxima.
std::vector<NodeDef> result;
result.push_back(NodeDef("Shape", {I(0)}, {IA("X_shape")}));
result.push_back(NodeDef("Expand", {O(0), IA("X_shape")}, {IA("expanded_Y")}));
result.push_back(NodeDef("Equal", {I(0), IA("expanded_Y")}, {IA("mask")}));
result.push_back(NodeDef("Cast",
{IA("mask")},
{IA("mask_cast")},
{MakeAttribute("to", static_cast<int64_t>(IElemType(0)))}));
result.push_back(NodeDef("Expand", {GO(0), IA("X_shape")}, {IA("expanded_dY")}));
result.push_back(NodeDef("Mul", {IA("mask_cast"), IA("expanded_dY")}, {GI(0)}));
return result;
}
} // namespace training
} // namespace onnxruntime

View file

@ -94,6 +94,7 @@ DECLARE_GRADIENT_BUILDER(GetLeakyReluGradient)
DECLARE_GRADIENT_BUILDER(GetConvTransposeGradient)
DECLARE_GRADIENT_BUILDER(GetResizeGradient)
DECLARE_GRADIENT_BUILDER(GetAtanGradient)
DECLARE_GRADIENT_BUILDER(GetGlobalMaxPoolGradient)
DECLARE_GRADIENT_BUILDER(GetExternalGradient)

View file

@ -126,6 +126,7 @@ void GradientBuilderRegistry::RegisterGradientBuilders() {
REGISTER_GRADIENT_BUILDER("ConvTranspose", GetConvTransposeGradient);
REGISTER_GRADIENT_BUILDER("Resize", GetResizeGradient);
REGISTER_GRADIENT_BUILDER("Atan", GetAtanGradient);
REGISTER_GRADIENT_BUILDER("GlobalMaxPool", GetGlobalMaxPoolGradient);
REGISTER_GRADIENT_BUILDER("ExternalGradient", GetExternalGradient);
};

View file

@ -3356,6 +3356,29 @@ TEST(GradientCheckerTest, ResizeGrad) {
TEST(GradientCheckerTest, AtanGrad) { UnaryOpGradientTest("Atan"); }
TEST(GradientCheckerTest, GlobalMaxPoolGrad) {
float max_error;
GradientChecker<float, float, float> gradient_checker;
OpDef op_def{"GlobalMaxPool", kOnnxDomain, 11};
constexpr float error_tolerance = 1e-3f;
// globalmaxpool
{
ASSERT_STATUS_OK(gradient_checker.ComputeGradientError(op_def, {{2, 3, 5, 5}}, {{2, 3, 1, 1}}, &max_error, {},
/*check_not_have_gradient*/ true,
/*check_not_have_shape_inferencing*/ true));
EXPECT_IS_TINIER_THAN(max_error, error_tolerance);
}
// globalmaxpool_precomputed
{
ASSERT_STATUS_OK(gradient_checker.ComputeGradientError(op_def, {{2, 1, 3, 3}}, {{2, 1, 1, 1}}, &max_error, {},
/*check_not_have_gradient*/ true,
/*check_not_have_shape_inferencing*/ true));
EXPECT_IS_TINIER_THAN(max_error, error_tolerance);
}
}
} // namespace test
} // namespace onnxruntime