diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index 43835f07c4..1ba52ca9e5 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -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 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(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 diff --git a/orttraining/orttraining/core/graph/gradient_builder.h b/orttraining/orttraining/core/graph/gradient_builder.h index 2b40754b62..2611e742f3 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.h +++ b/orttraining/orttraining/core/graph/gradient_builder.h @@ -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) diff --git a/orttraining/orttraining/core/graph/gradient_builder_registry.cc b/orttraining/orttraining/core/graph/gradient_builder_registry.cc index 9c9884c5d3..a04d909267 100755 --- a/orttraining/orttraining/core/graph/gradient_builder_registry.cc +++ b/orttraining/orttraining/core/graph/gradient_builder_registry.cc @@ -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); }; diff --git a/orttraining/orttraining/test/gradient/gradient_ops_test.cc b/orttraining/orttraining/test/gradient/gradient_ops_test.cc index b683010a72..f4083d5b8f 100644 --- a/orttraining/orttraining/test/gradient/gradient_ops_test.cc +++ b/orttraining/orttraining/test/gradient/gradient_ops_test.cc @@ -3356,6 +3356,29 @@ TEST(GradientCheckerTest, ResizeGrad) { TEST(GradientCheckerTest, AtanGrad) { UnaryOpGradientTest("Atan"); } +TEST(GradientCheckerTest, GlobalMaxPoolGrad) { + float max_error; + GradientChecker 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