From a5d3a52d1a8f09560982497bb5482f58559e705d Mon Sep 17 00:00:00 2001 From: harshithapv <54084812+harshithapv@users.noreply.github.com> Date: Tue, 13 Apr 2021 12:54:45 -0700 Subject: [PATCH] Add Tile grad (#7289) * tile grad * fixed bugs * added tile grad test * bug fix * Added tests. Addressed comments * added optimization recommended and addressed comments * fixed comment --- .../core/graph/gradient_builder.cc | 103 ++++++++++++++---- .../orttraining/core/graph/gradient_builder.h | 1 + .../core/graph/gradient_builder_registry.cc | 1 + .../test/gradient/gradient_ops_test.cc | 54 +++++++++ 4 files changed, 140 insertions(+), 19 deletions(-) mode change 100644 => 100755 orttraining/orttraining/core/graph/gradient_builder.cc mode change 100644 => 100755 orttraining/orttraining/core/graph/gradient_builder.h mode change 100644 => 100755 orttraining/orttraining/core/graph/gradient_builder_registry.cc mode change 100644 => 100755 orttraining/orttraining/test/gradient/gradient_ops_test.cc diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc old mode 100644 new mode 100755 index 011890f6ad..1aca9b6342 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -749,7 +749,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetUnsqueezeGradient) { {GO(0)}, {GI(0)}, SrcNodeAttributes())}; - } else { // mandatory input 'axes' since opset 13 + } else { // mandatory input 'axes' since opset 13 return std::vector{ NodeDef(OpDef{"Squeeze", kOnnxDomain, 13}, {GO(0), I(1)}, @@ -790,7 +790,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetReluGradient) { IMPLEMENT_GRADIENT_BUILDER(GetSqueezeGradient) { std::vector result; size_t numInputs = GetSrcNodeInputSize(); - if (SrcNodeOpsetVersion() < 13) { //axes attribute + if (SrcNodeOpsetVersion() < 13) { //axes attribute auto attributes = SrcNodeAttributes(); std::vector axes_values; if (attributes.find("axes") != attributes.end()) { @@ -800,21 +800,21 @@ IMPLEMENT_GRADIENT_BUILDER(GetSqueezeGradient) { {GO(0)}, {GI(0)}, {MakeAttribute("axes", axes_values)})); - } - } else if(numInputs == 2){ //optional input 'axes' is provided + } + } else if (numInputs == 2) { //optional input 'axes' is provided result.push_back( NodeDef(OpDef{"Unsqueeze", kOnnxDomain, 13}, {GO(0), I(1)}, {GI(0)})); - } else { // if axes attribute/input not provided for squeeze - result.push_back( - NodeDef("Shape", - {I(0)}, - {IA("I0_shape")})); - result.push_back( - NodeDef("Reshape", - {GO(0), IA("I0_shape")}, - {GI(0)})); + } else { // if axes attribute/input not provided for squeeze + result.push_back( + NodeDef("Shape", + {I(0)}, + {IA("I0_shape")})); + result.push_back( + NodeDef("Reshape", + {GO(0), IA("I0_shape")}, + {GI(0)})); } return result; @@ -1123,12 +1123,11 @@ IMPLEMENT_GRADIENT_BUILDER(GetReduceSumGradient) { grad = IA("Unqueezed_Grad"); result.push_back(NodeDef("Unsqueeze", {GO(0)}, {grad}, {MakeAttribute("axes", axes_values)})); - - } + } } else if (numInputs == 2) { //optional input 'axes' is available as input I(1) grad = IA("Unqueezed_Grad"); result.push_back(NodeDef(OpDef{"Unsqueeze", kOnnxDomain, 13}, {GO(0), I(1)}, {grad})); - } //axes is not available, the GO(0) is a scalar which can be expanded to required shape + } //axes is not available, the GO(0) is a scalar which can be expanded to required shape } result.push_back(NodeDef("Shape", {I(0)}, {IA("Shaped_X")})); @@ -1259,7 +1258,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetFastGeluGradient) { ArgDef x_shape = IA("Shape_" + X.name); return GetBiasGeluGradNodes(true, dY, X, B, dX, dB, b_axes, b_shape, x_shape, NodeName()); } - + if (num_src_node_inputs == 1) { // without bias return std::vector{ NodeDef(OpDef{"FastGeluGrad", kMSDomain, 1}, @@ -1484,8 +1483,74 @@ IMPLEMENT_GRADIENT_BUILDER(GetClipGradient) { IMPLEMENT_GRADIENT_BUILDER(GetAbsGradient) { return std::vector{ NodeDef("Sign", {I(0)}, {IA("Sign_Input")}), - NodeDef("Mul", {GO(0), IA("Sign_Input")}, {GI(0)}) - }; + NodeDef("Mul", {GO(0), IA("Sign_Input")}, {GI(0)})}; +} + +// Computes gradient of Tile Operation. +// 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] +// 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)) +// then perform reducesum on the reshaped Y^_grad on its even indices to get X_grad. +// even_indices = [0, 2, 4...] +// X_grad = ReduceSum(Y^_grad, even_indices) + +IMPLEMENT_GRADIENT_BUILDER(GetTileGradient) { + std::vector result = {}; + + result.push_back(NodeDef("Shape", {I(0)}, {IA("orig_shape")})); + std::vector axes_values = {1}; + result.push_back(NodeDef("Unsqueeze", {IA("orig_shape")}, {IA("2d_orig_shape")}, {MakeAttribute("axes", axes_values)})); // M, N, K + result.push_back(NodeDef("Unsqueeze", {I(1)}, {IA("2d_repeats")}, {MakeAttribute("axes", axes_values)})); //a, b, c + result.push_back(NodeDef("Concat", {IA("2d_repeats"), IA("2d_orig_shape")}, {IA("concated_dims_T")}, + {MakeAttribute("axis", int64_t(1))})); // [[a, M], [b, N], [c, K]] + std::vector const_shape_minusone{-1}; + NodeDef const_shape_minusone_node = ConstantVectorNode(const_shape_minusone, Name("const_shape_minusone")); + result.push_back(const_shape_minusone_node); + result.push_back(NodeDef("Reshape", {IA("concated_dims_T"), const_shape_minusone_node.output_args[0]}, + {IA("concated_dims_flatten")})); // flatten [a, M, b, N, c, K] + + result.push_back(NodeDef("Reshape", {GO(0), IA("concated_dims_flatten")}, {IA("reshape_tile_grad_op")})); + + std::vector orig_shape, repeat_shape; + bool orig_has_shape = GetShape(I(0), orig_shape).IsOK(); + bool repeat_has_shape = GetShape(I(1), repeat_shape).IsOK(); + + if (orig_has_shape || repeat_has_shape) { + int64_t limit = orig_has_shape ? orig_shape.size() : repeat_shape[0].dim_value(); + limit = 2 * limit; + + std::vector even_indices; + + for (int64_t i = 0; i < limit; i = i + 2) { + even_indices.push_back(i); + } + NodeDef even_indices_node = ConstantVectorNode(even_indices, Name("even_indices")); + result.push_back(even_indices_node); + int opset_version = SrcNodeDomain() == kOnnxDomain ? SrcNodeOpsetVersion() : OnnxOpSetVersion(); + result.push_back(NodeDef(opset_version >= 13 ? OpDef{"ReduceSum", kOnnxDomain, opset_version} : OpDef{"ReduceSumTraining", kMSDomain, 1}, + {IA("reshape_tile_grad_op"), even_indices_node.output_args[0]}, + {GI(0)}, + {{"keepdims", ONNX_NAMESPACE::MakeAttribute("keepdims", int64_t{0})}})); + + } else { + NodeDef start_node = ConstantScalarNode(int64_t{0}, {}, Name("start_int64")); + NodeDef delta_node = ConstantScalarNode(int64_t{2}, {}, Name("delta_int64")); + result.push_back(NodeDef("Size", {IA("concated_dims_flatten")}, {IA("limit")})); // get num dimensions of the flattened grad op = 6 + result.push_back(start_node); + result.push_back(delta_node); + result.push_back(NodeDef("Range", {start_node.output_args[0], IA("limit"), delta_node.output_args[0]}, {IA("range_even_indices")})); + + int opset_version = SrcNodeDomain() == kOnnxDomain ? SrcNodeOpsetVersion() : OnnxOpSetVersion(); + result.push_back(NodeDef(opset_version >= 13 ? OpDef{"ReduceSum", kOnnxDomain, opset_version} : OpDef{"ReduceSumTraining", kMSDomain, 1}, + {IA("reshape_tile_grad_op"), IA("range_even_indices")}, + {GI(0)}, + {{"keepdims", ONNX_NAMESPACE::MakeAttribute("keepdims", int64_t{0})}})); + } + return result; } } // namespace training diff --git a/orttraining/orttraining/core/graph/gradient_builder.h b/orttraining/orttraining/core/graph/gradient_builder.h old mode 100644 new mode 100755 index d8f336c5b5..fc5df7ab53 --- 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(GetTileGradient) } // namespace training } // namespace onnxruntime diff --git a/orttraining/orttraining/core/graph/gradient_builder_registry.cc b/orttraining/orttraining/core/graph/gradient_builder_registry.cc old mode 100644 new mode 100755 index 07865853df..e4430f36e4 --- a/orttraining/orttraining/core/graph/gradient_builder_registry.cc +++ b/orttraining/orttraining/core/graph/gradient_builder_registry.cc @@ -101,6 +101,7 @@ void GradientBuilderRegistry::RegisterGradientBuilders() { REGISTER_GRADIENT_BUILDER("TopK", GetTopKGradient); REGISTER_GRADIENT_BUILDER("Clip", GetClipGradient); REGISTER_GRADIENT_BUILDER("Abs", GetAbsGradient); + REGISTER_GRADIENT_BUILDER("Tile", GetTileGradient); }; } // namespace training diff --git a/orttraining/orttraining/test/gradient/gradient_ops_test.cc b/orttraining/orttraining/test/gradient/gradient_ops_test.cc old mode 100644 new mode 100755 index 917f401933..114ce159f5 --- a/orttraining/orttraining/test/gradient/gradient_ops_test.cc +++ b/orttraining/orttraining/test/gradient/gradient_ops_test.cc @@ -2437,6 +2437,60 @@ TEST(GradientCheckerTest, ClipGrad) { } } +TEST(GradientCheckerTest, TileGrad) { + float max_error; + GradientChecker gradient_checker; + OpDef op_def{"Tile", kOnnxDomain, 11}; + + // 2D input + { + TensorInfo x_info({2, 4}, true); + TensorInfo repeat_info({2}, false, nullptr, DataTypeImpl::GetTensorType()); + std::vector> x_datas = {{1, 2, 3, 4, 5, 6, 7, 8}, {2, 2}}; + + TensorInfo y_info({4, 8}, true); + + gradient_checker.ComputeGradientError(op_def, {x_info, repeat_info}, {y_info}, &max_error, x_datas); + EXPECT_IS_TINY(max_error); + } + + // 1D input + { + TensorInfo x_info({2}, true); + TensorInfo repeat_info({1}, false, nullptr, DataTypeImpl::GetTensorType()); + std::vector> x_datas = {{1, 2}, {4}}; + + TensorInfo y_info({8}, true); + + gradient_checker.ComputeGradientError(op_def, {x_info, repeat_info}, {y_info}, &max_error, x_datas); + EXPECT_IS_TINY(max_error); + } + + // 3D input + { + TensorInfo x_info({2, 2, 3}, true); + TensorInfo repeat_info({3}, false, nullptr, DataTypeImpl::GetTensorType()); + std::vector> x_datas = {{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, {2, 3, 4}}; + + TensorInfo y_info({4, 6, 12}, true); + + gradient_checker.ComputeGradientError(op_def, {x_info, repeat_info}, {y_info}, &max_error, x_datas); + EXPECT_IS_TINY(max_error); + } + + // 3D input - repeating 1s + { + TensorInfo x_info({2, 2, 3}, true); + TensorInfo repeat_info({3}, false, nullptr, DataTypeImpl::GetTensorType()); + std::vector> x_datas = {{1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}, {1, 1, 1}}; + + TensorInfo y_info({2, 2, 3}, true); + + gradient_checker.ComputeGradientError(op_def, {x_info, repeat_info}, {y_info}, &max_error, x_datas); + EXPECT_IS_TINY(max_error); + } +} + } // namespace test } // namespace onnxruntime