From 1a3ddf0714e1bdf9b807a342eee5f6e160ad1ec9 Mon Sep 17 00:00:00 2001 From: mindest <30493312+mindest@users.noreply.github.com> Date: Tue, 20 Apr 2021 18:14:31 +0800 Subject: [PATCH] Add gradient registration and tests for Min/Max (#7217) * Add gradient registration and tests for Min/Max * Add helper function for min/max grad test * limit Min/Max Grad to accept at most two inputs; modify test case accordingly * resolve merge error --- .../core/graph/gradient_builder.cc | 60 ++++++++++++++++++- .../orttraining/core/graph/gradient_builder.h | 1 + .../core/graph/gradient_builder_registry.cc | 2 + .../test/gradient/gradient_ops_test.cc | 38 ++++++++++++ 4 files changed, 100 insertions(+), 1 deletion(-) diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index 1aca9b6342..9630f45125 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -1490,7 +1490,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetAbsGradient) { // Tile is defined as follows: // Y = Tile(X, repeat), say, // X shape : M, N, K -// repeat is a 1D tensor with value: [a, b, c] +// repeat is a 1D tensor with value: [a, b, c] // Y shape : aM, bN, cK // To compute the gradient of y, we first reshape the gradient of y as, // Y^_grad = Reshape(Y_grad(a, M, b, N, c, K)) @@ -1553,5 +1553,63 @@ IMPLEMENT_GRADIENT_BUILDER(GetTileGradient) { return result; } +IMPLEMENT_GRADIENT_BUILDER(GetMinMaxGradient) { + const auto num_src_node_inputs = GetSrcNodeInputSize(); + if (num_src_node_inputs == 1) { + if (IsGradientRequiredForSrcNodeInput(0)) { + return std::vector{NodeDef("Identity", {GO(0)}, {GI(0)})}; + } + + return std::vector{}; + } + + if (num_src_node_inputs > 2) { + ORT_THROW("Min/Max gradient currently does not support over 2 inputs."); + } + + if (!IsGradientRequiredForSrcNodeInput(0) && !IsGradientRequiredForSrcNodeInput(1)) { + return std::vector{}; + } + + std::vector result; + std::vector y_shape; + const ArgDef y = O(0); + bool get_y_shape_ok = GetShape(y, y_shape).IsOK(); + result.push_back(NodeDef("Equal", {I(1), y}, {IA("Mask_1")})); + if (IsGradientRequiredForSrcNodeInput(0)) { + result.push_back(NodeDef("Not", {IA("Mask_1")}, {IA("Mask_0")})); + } + for (int i = 0; i < num_src_node_inputs; i++) { + if (IsGradientRequiredForSrcNodeInput(i)) { + const ArgDef x = I(i); + const ArgDef mask_cast_i_def = IA("Mask_Cast_" + std::to_string(i)); + const ArgDef pre_reduce_grad_i_def = IA("PreReduceGrad_" + std::to_string(i), OType(0)); + result.push_back(NodeDef("Cast", + {IA("Mask_" + std::to_string(i))}, + {mask_cast_i_def}, + {MakeAttribute("to", int64_t(IElemType(0)))})); + result.push_back(NodeDef("Mul", {mask_cast_i_def, GO(0)}, {pre_reduce_grad_i_def})); + std::vector x_shape; + if (get_y_shape_ok && GetShape(x, x_shape).IsOK()) { + std::vector x_axes; + ComputeBroadcastBackwardAxes(x_shape, y_shape, &x_axes, nullptr, NodeName()); + if (x_axes.size() > 0) { + HandleBroadcasting(pre_reduce_grad_i_def, x, GI(i), x_axes, result); + } else { + result.push_back(NodeDef("Identity", {pre_reduce_grad_i_def}, {GI(i)})); + } + } else { + ArgDef x_axes_def = IA("ReduceAxes_" + x.name); + ArgDef x_shape_def = IA("Shape_" + x.name); + ArgDef y_shape_def = IA("Shape_" + y.name + std::to_string(i)); + ComputeBroadcastBackwardAxesDynamic(x, y, x_shape_def, y_shape_def, &x_axes_def, nullptr, result); + HandleBroadcastingDynamic(pre_reduce_grad_i_def, x, x_shape_def, GI(i), x_axes_def, result); + } + } + } + + 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 fc5df7ab53..1a1a0ab002 100755 --- a/orttraining/orttraining/core/graph/gradient_builder.h +++ b/orttraining/orttraining/core/graph/gradient_builder.h @@ -70,6 +70,7 @@ DECLARE_GRADIENT_BUILDER(GetFlattenGradient) DECLARE_GRADIENT_BUILDER(GetTopKGradient) DECLARE_GRADIENT_BUILDER(GetClipGradient) DECLARE_GRADIENT_BUILDER(GetAbsGradient) +DECLARE_GRADIENT_BUILDER(GetMinMaxGradient) DECLARE_GRADIENT_BUILDER(GetTileGradient) } // namespace training diff --git a/orttraining/orttraining/core/graph/gradient_builder_registry.cc b/orttraining/orttraining/core/graph/gradient_builder_registry.cc index e4430f36e4..7399f96308 100755 --- a/orttraining/orttraining/core/graph/gradient_builder_registry.cc +++ b/orttraining/orttraining/core/graph/gradient_builder_registry.cc @@ -101,6 +101,8 @@ void GradientBuilderRegistry::RegisterGradientBuilders() { REGISTER_GRADIENT_BUILDER("TopK", GetTopKGradient); REGISTER_GRADIENT_BUILDER("Clip", GetClipGradient); REGISTER_GRADIENT_BUILDER("Abs", GetAbsGradient); + REGISTER_GRADIENT_BUILDER("Min", GetMinMaxGradient); + REGISTER_GRADIENT_BUILDER("Max", GetMinMaxGradient); REGISTER_GRADIENT_BUILDER("Tile", GetTileGradient); }; diff --git a/orttraining/orttraining/test/gradient/gradient_ops_test.cc b/orttraining/orttraining/test/gradient/gradient_ops_test.cc index 1ea29d8110..67b2da8616 100755 --- a/orttraining/orttraining/test/gradient/gradient_ops_test.cc +++ b/orttraining/orttraining/test/gradient/gradient_ops_test.cc @@ -2439,6 +2439,44 @@ TEST(GradientCheckerTest, ClipGrad) { } } +void GradientCheckerMinMaxGradHelper(const std::string op) { + float max_error; + GradientChecker gradient_checker; + OpDef op_def{op, kOnnxDomain, 11}; + + // Exclude equal inputs, since Min/Max is not smooth in such case + { + TensorInfo x_info({2, 3}, true); + TensorInfo y_info({2, 3}, true); + gradient_checker.ComputeGradientError(op_def, {x_info}, {y_info}, &max_error); + EXPECT_IS_TINY(max_error); + } + + { + TensorInfo x1_info({2, 3}, true); + TensorInfo x2_info({2, 3}, true); + TensorInfo y_info({2, 3}, true); + gradient_checker.ComputeGradientError(op_def, {x1_info, x2_info}, {y_info}, &max_error); + EXPECT_IS_TINY(max_error); + } + + { + TensorInfo x1_info({2, 3}, true); + TensorInfo x2_info({3}, true); + TensorInfo y_info({2, 3}, true); + gradient_checker.ComputeGradientError(op_def, {x1_info, x2_info}, {y_info}, &max_error); + EXPECT_IS_TINY(max_error); + } +} + +TEST(GradientCheckerTest, MinGrad) { + GradientCheckerMinMaxGradHelper("Min"); +} + +TEST(GradientCheckerTest, MaxGrad) { + GradientCheckerMinMaxGradHelper("Max"); +} + TEST(GradientCheckerTest, TileGrad) { float max_error; GradientChecker gradient_checker;