From 03996c7c0846020dbffe8daad93d801ceb22b5e5 Mon Sep 17 00:00:00 2001 From: Vincent Wang Date: Fri, 10 Apr 2020 19:35:32 +0800 Subject: [PATCH] Fixes for Where, ConcatGrad and ReduceSumGrad (#3415) * Fixes for Expand, Where, ConcatGrad ReduceSumGrad. * Roll back expand, fix, add tests for reduce grad. * Roll back CPU Expand change. * Fix after merge. Co-authored-by: Vincent Wang --- onnxruntime/core/providers/cpu/tensor/utils.h | 2 +- .../core/providers/cuda/tensor/where.cc | 42 ++++--- .../core/graph/gradient_builder.cc | 83 ++++++++++++-- .../orttraining/core/graph/gradient_builder.h | 1 + .../core/graph/gradient_builder_registry.cc | 1 + .../test/gradient/gradient_ops_test.cc | 107 +++++++++++++++++- 6 files changed, 208 insertions(+), 28 deletions(-) diff --git a/onnxruntime/core/providers/cpu/tensor/utils.h b/onnxruntime/core/providers/cpu/tensor/utils.h index 5fd088030e..358449ba16 100644 --- a/onnxruntime/core/providers/cpu/tensor/utils.h +++ b/onnxruntime/core/providers/cpu/tensor/utils.h @@ -34,7 +34,7 @@ struct TensorPitches : std::vector { if (padded_rank >= 1) { for (size_t i = 0; i < padded_rank; ++i) { - if (i == 0) + if (i == 0 && tensor_rank > 0) // For scalar tensor, the values in the pitches are all 1. p.operator[](padded_rank - 1) = p.operator[](padded_rank) * dims[0]; else p.operator[](padded_rank - 1 - i) = p.operator[](padded_rank - 1); diff --git a/onnxruntime/core/providers/cuda/tensor/where.cc b/onnxruntime/core/providers/cuda/tensor/where.cc index 3fc0b9272e..b3e46c3983 100644 --- a/onnxruntime/core/providers/cuda/tensor/where.cc +++ b/onnxruntime/core/providers/cuda/tensor/where.cc @@ -95,37 +95,43 @@ struct TernaryElementwisePreparation { output_rank_or_simple_broadcast = out_rank; if (a_shape != output_shape) { - TensorPitches a_pitches(a_shape.GetDims()); a_padded_strides.size_ = out_rank; - auto offset = out_rank - a_rank; - for (auto i = offset; i < out_rank; ++i) { - // the stride for broadcast dimension is kept as 0 - if (a_shape.GetDims()[i - offset] != 1) { - a_padded_strides[i] = a_pitches[i]; + if (a_rank > 0) { + TensorPitches a_pitches(a_shape.GetDims()); + auto offset = out_rank - a_rank; + for (auto i = offset; i < out_rank; ++i) { + // the stride for broadcast dimension is kept as 0 + if (a_shape.GetDims()[i - offset] != 1) { + a_padded_strides[i] = a_pitches[i]; + } } } } if (b_shape != output_shape) { - TensorPitches b_pitches(b_shape.GetDims()); b_padded_strides.size_ = out_rank; - auto offset = out_rank - b_rank; - for (auto i = offset; i < out_rank; ++i) { - // the stride for broadcast dimension is kept as 0 - if (b_shape.GetDims()[i - offset] != 1) { - b_padded_strides[i] = b_pitches[i]; + if (b_rank > 0) { + TensorPitches b_pitches(b_shape.GetDims()); + auto offset = out_rank - b_rank; + for (auto i = offset; i < out_rank; ++i) { + // the stride for broadcast dimension is kept as 0 + if (b_shape.GetDims()[i - offset] != 1) { + b_padded_strides[i] = b_pitches[i]; + } } } } if (c_shape != output_shape) { - TensorPitches c_pitches(c_shape.GetDims()); c_padded_strides.size_ = out_rank; - auto offset = out_rank - c_rank; - for (auto i = offset; i < out_rank; ++i) { - // the stride for broadcast dimension is kept as 0 - if (c_shape.GetDims()[i - offset] != 1) { - c_padded_strides[i] = c_pitches[i]; + if (c_rank > 0) { + TensorPitches c_pitches(c_shape.GetDims()); + auto offset = out_rank - c_rank; + for (auto i = offset; i < out_rank; ++i) { + // the stride for broadcast dimension is kept as 0 + if (c_shape.GetDims()[i - offset] != 1) { + c_padded_strides[i] = c_pitches[i]; + } } } } diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index 1a76f0dc2b..82e49e9f68 100644 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -408,18 +408,32 @@ IMPLEMENT_GRADIENT_BUILDER(GetSplitGradient) { } IMPLEMENT_GRADIENT_BUILDER(GetConcatGradient) { - //TODO: split attribute should be used!!! - //AttributeProto split = MakeAttribute("split", std::vector()); + auto attributes = SrcNodeAttributes(); + ORT_ENFORCE(attributes.at("axis").has_i()); + auto axis = attributes.at("axis").i(); + std::vector split_attribute(GetSrcNodeInputSize()); std::vector outputs; for (int i = 0; i < GetSrcNodeInputSize(); ++i) { + std::vector data_shape = GetShape(I(i)); + int64_t axis_index = axis < 0 ? static_cast(data_shape.size()) + axis : axis; + if (axis_index >= 0 && axis_index < static_cast(data_shape.size()) && data_shape[axis_index].has_dim_value()) { + split_attribute[i] = data_shape[axis_index].dim_value(); + } else { + ORT_THROW("Error: can't infer split attribute value for ConcatGrad"); + } outputs.push_back(GI(i)); } + + std::vector new_attributes; + new_attributes.push_back(MakeAttribute("axis", axis)); + new_attributes.push_back(MakeAttribute("split", split_attribute)); + return std::vector{ NodeDef("Split", {GO(0)}, outputs, - SrcNodeAttributes())}; + new_attributes)}; } IMPLEMENT_GRADIENT_BUILDER(GetGatherNDGradient) { @@ -728,6 +742,11 @@ IMPLEMENT_GRADIENT_BUILDER(GetReduceMeanGradient) { std::vector axes_values(data_shape.size()); if (attributes.find("axes") != attributes.end()) { axes_values = RetrieveValues(attributes.at("axes")); + for (size_t i = 0; i < axes_values.size(); i++) { + if (axes_values[i] < 0) { + axes_values[i] = data_shape.size() + axes_values[i]; + } + } } else { std::iota(std::begin(axes_values), std::end(axes_values), 0); } @@ -751,10 +770,6 @@ IMPLEMENT_GRADIENT_BUILDER(GetReduceMeanGradient) { std::vector repeats(data_shape.size(), 1); int64_t scale = 1; for (int64_t axis : axes_values) { - if (axis < 0) { - axis = data_shape.size() + axis; - } - if (data_shape[axis].has_dim_value()) { auto dim_value = data_shape[axis].dim_value(); repeats[axis] = dim_value; @@ -783,6 +798,60 @@ IMPLEMENT_GRADIENT_BUILDER(GetReduceMeanGradient) { return result; } +IMPLEMENT_GRADIENT_BUILDER(GetReduceSumGradient) { + std::vector data_shape = GetShape(I(0)); + std::vector result; + + auto attributes = SrcNodeAttributes(); + std::vector axes_values(data_shape.size()); + if (attributes.find("axes") != attributes.end()) { + axes_values = RetrieveValues(attributes.at("axes")); + for (size_t i = 0; i < axes_values.size(); i++) { + if (axes_values[i] < 0) { + axes_values[i] = data_shape.size() + axes_values[i]; + } + } + } else { + std::iota(std::begin(axes_values), std::end(axes_values), 0); + } + + bool keepdims = true; + if (attributes.find("keepdims") != attributes.end() && + attributes.at("keepdims").has_i()) { + keepdims = static_cast(attributes.at("keepdims").i()); + } + + ArgDef unsqueezed_Grad = GO(0); + if (!keepdims) { + unsqueezed_Grad = IA("Unqueezed_Grad"); + result.push_back( + NodeDef("Unsqueeze", + {GO(0)}, + {unsqueezed_Grad}, + {MakeAttribute("axes", axes_values)})); + } + + std::vector repeats(data_shape.size(), 1); + for (int64_t axis : axes_values) { + if (data_shape[axis].has_dim_value()) { + auto dim_value = data_shape[axis].dim_value(); + repeats[axis] = dim_value; + } else { + ORT_THROW("Error: can't infer repeats value for ReduceSumGrad"); + } + } + + NodeDef repeats_node = ConstantValueNode(repeats, Name("repeats")); + ArgDef REPEATS = repeats_node.output_args[0]; + result.push_back(repeats_node); + result.push_back( + NodeDef("Tile", + {unsqueezed_Grad, REPEATS}, + {GI(0)})); + + return result; +} + IMPLEMENT_GRADIENT_BUILDER(GetPowGradient) { if (IsGradientRequiredForSrcNodeInput(1)) { ORT_THROW("GradientBuilder is not implemented for CUDA Pow's input exponent."); diff --git a/orttraining/orttraining/core/graph/gradient_builder.h b/orttraining/orttraining/core/graph/gradient_builder.h index 3de555b3de..cbbc25ad5d 100644 --- a/orttraining/orttraining/core/graph/gradient_builder.h +++ b/orttraining/orttraining/core/graph/gradient_builder.h @@ -24,6 +24,7 @@ DECLARE_GRADIENT_BUILDER(GetAddSubGradient) DECLARE_GRADIENT_BUILDER(GetMulGradient) DECLARE_GRADIENT_BUILDER(GetDivGradient) DECLARE_GRADIENT_BUILDER(GetReduceMeanGradient) +DECLARE_GRADIENT_BUILDER(GetReduceSumGradient) DECLARE_GRADIENT_BUILDER(GetPowGradient) DECLARE_GRADIENT_BUILDER(GetConcatGradient) DECLARE_GRADIENT_BUILDER(GetReshapeGradient) diff --git a/orttraining/orttraining/core/graph/gradient_builder_registry.cc b/orttraining/orttraining/core/graph/gradient_builder_registry.cc index 057e896ebc..809d5dc2f8 100644 --- a/orttraining/orttraining/core/graph/gradient_builder_registry.cc +++ b/orttraining/orttraining/core/graph/gradient_builder_registry.cc @@ -60,6 +60,7 @@ void GradientBuilderRegistry::RegisterGradientBuilders() { REGISTER_GRADIENT_BUILDER("Relu", GetReluGradient); REGISTER_GRADIENT_BUILDER("Pow", GetPowGradient); REGISTER_GRADIENT_BUILDER("ReduceMean", GetReduceMeanGradient); + REGISTER_GRADIENT_BUILDER("ReduceSum", GetReduceSumGradient); REGISTER_GRADIENT_BUILDER("Add", GetAddSubGradient); REGISTER_GRADIENT_BUILDER("Sub", GetAddSubGradient); REGISTER_GRADIENT_BUILDER("Mul", GetMulGradient); diff --git a/orttraining/orttraining/test/gradient/gradient_ops_test.cc b/orttraining/orttraining/test/gradient/gradient_ops_test.cc index b7a4b7ce8a..176ab954d4 100644 --- a/orttraining/orttraining/test/gradient/gradient_ops_test.cc +++ b/orttraining/orttraining/test/gradient/gradient_ops_test.cc @@ -433,6 +433,89 @@ TEST(GradientCheckerTest, ReduceMeanGrad) { MakeAttribute("keepdims", int64_t(0))}); EXPECT_IS_TINY(max_error); } + + // axes = [-2], keepdims = 1 + { + gradient_checker.ComputeGradientError(op_def, {{4, 3, 2}}, {{4, 1, 2}}, &max_error, + {MakeAttribute("axes", std::vector{-2}), + MakeAttribute("keepdims", int64_t(1))}); + EXPECT_IS_TINY(max_error); + } + + // axes = [-2, -1], keepdims = 0 + { + gradient_checker.ComputeGradientError(op_def, {{4, 3, 2}}, {{4}}, &max_error, + {MakeAttribute("axes", std::vector{-2, -1}), + MakeAttribute("keepdims", int64_t(0))}); + EXPECT_IS_TINY(max_error); + } +} + +TEST(GradientCheckerTest, ReduceSumGrad) { + float max_error; + GradientChecker gradient_checker; + OpDef op_def{"ReduceSum"}; + + // default + { + gradient_checker.ComputeGradientError(op_def, {{4, 3, 2}}, {{1, 1, 1}}, &max_error); + EXPECT_IS_TINY(max_error); + } + + // axes = [0, 1, 2], keepdims = 0 + { + gradient_checker.ComputeGradientError(op_def, {{4, 3, 2}}, {{}}, &max_error, + {MakeAttribute("axes", std::vector{0, 1, 2}), + MakeAttribute("keepdims", int64_t(0))}); + EXPECT_IS_TINY(max_error); + } + + // axes = [0, 2], keepdims = 1 + { + gradient_checker.ComputeGradientError(op_def, {{4, 3, 2}}, {{1, 3, 1}}, &max_error, + {MakeAttribute("axes", std::vector{0, 2})}); + EXPECT_IS_TINY(max_error); + } + + // axes = [0, 1], keepdims = 0 + { + gradient_checker.ComputeGradientError(op_def, {{4, 3, 2}}, {{2}}, &max_error, + {MakeAttribute("axes", std::vector{0, 1}), + MakeAttribute("keepdims", int64_t(0))}); + EXPECT_IS_TINY(max_error); + } + + // axes = [1], keepdims = 1 + { + gradient_checker.ComputeGradientError(op_def, {{4, 3, 2}}, {{4, 1, 2}}, &max_error, + {MakeAttribute("axes", std::vector{1}), + MakeAttribute("keepdims", int64_t(1))}); + EXPECT_IS_TINY(max_error); + } + + // axes = [2], keepdims = 0 + { + gradient_checker.ComputeGradientError(op_def, {{4, 3, 2}}, {{4, 3}}, &max_error, + {MakeAttribute("axes", std::vector{2}), + MakeAttribute("keepdims", int64_t(0))}); + EXPECT_IS_TINY(max_error); + } + + // axes = [-2], keepdims = 1 + { + gradient_checker.ComputeGradientError(op_def, {{4, 3, 2}}, {{4, 1, 2}}, &max_error, + {MakeAttribute("axes", std::vector{-2}), + MakeAttribute("keepdims", int64_t(1))}); + EXPECT_IS_TINY(max_error); + } + + // axes = [-1, -3], keepdims = 0 + { + gradient_checker.ComputeGradientError(op_def, {{4, 3, 2}}, {{3}}, &max_error, + {MakeAttribute("axes", std::vector{-1, -3}), + MakeAttribute("keepdims", int64_t(0))}); + EXPECT_IS_TINY(max_error); + } } #ifndef USE_CUDA @@ -626,6 +709,26 @@ TEST(GradientCheckerTest, ConcatGrad) { {MakeAttribute("axis", int64_t(2))}); EXPECT_IS_TINY(max_error); } + + //concat_different_shape + { + TensorShape x1_shape({2, 2}); + TensorShape x2_shape({2, 4}); + TensorShape y_shape({2, 6}); + gradient_checker.ComputeGradientError(op_def, {x1_shape, x2_shape}, {y_shape}, &max_error, + {MakeAttribute("axis", int64_t(1))}); + EXPECT_IS_TINY(max_error); + } + + //concat_different_shape_and_negative_axis + { + TensorShape x1_shape({2, 2}); + TensorShape x2_shape({2, 4}); + TensorShape y_shape({2, 6}); + gradient_checker.ComputeGradientError(op_def, {x1_shape, x2_shape}, {y_shape}, &max_error, + {MakeAttribute("axis", int64_t(-1))}); + EXPECT_IS_TINY(max_error); + } } TEST(GradientCheckerTest, AveragePoolGrad) { @@ -1554,11 +1657,11 @@ TEST(Synchronization, WaitAndRecordEventMany) { const size_t event_count = 16; for (int i = 0; i < 8; ++i) { std::thread thread_pool[2 * event_count]; - for (int j = 0; j < event_count; ++j) { + for (int j = 0; j < static_cast(event_count); ++j) { thread_pool[j] = std::thread(wait_event, j); thread_pool[j + event_count] = std::thread(record_event, j); } - for (int j = 0; j < event_count; ++j) { + for (size_t j = 0; j < event_count; ++j) { thread_pool[j].join(); thread_pool[j + event_count].join(); }