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 <weicwang@microsoft.com>
This commit is contained in:
Vincent Wang 2020-04-10 19:35:32 +08:00 committed by GitHub
parent e7297e6c9d
commit 03996c7c08
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 208 additions and 28 deletions

View file

@ -34,7 +34,7 @@ struct TensorPitches : std::vector<int64_t> {
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);

View file

@ -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];
}
}
}
}

View file

@ -408,18 +408,32 @@ IMPLEMENT_GRADIENT_BUILDER(GetSplitGradient) {
}
IMPLEMENT_GRADIENT_BUILDER(GetConcatGradient) {
//TODO: split attribute should be used!!!
//AttributeProto split = MakeAttribute("split", std::vector<int64_t>());
auto attributes = SrcNodeAttributes();
ORT_ENFORCE(attributes.at("axis").has_i());
auto axis = attributes.at("axis").i();
std::vector<int64_t> split_attribute(GetSrcNodeInputSize());
std::vector<ArgDef> outputs;
for (int i = 0; i < GetSrcNodeInputSize(); ++i) {
std::vector<Dimension> data_shape = GetShape(I(i));
int64_t axis_index = axis < 0 ? static_cast<int64_t>(data_shape.size()) + axis : axis;
if (axis_index >= 0 && axis_index < static_cast<int64_t>(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<AttributeProto> new_attributes;
new_attributes.push_back(MakeAttribute("axis", axis));
new_attributes.push_back(MakeAttribute("split", split_attribute));
return std::vector<NodeDef>{
NodeDef("Split",
{GO(0)},
outputs,
SrcNodeAttributes())};
new_attributes)};
}
IMPLEMENT_GRADIENT_BUILDER(GetGatherNDGradient) {
@ -728,6 +742,11 @@ IMPLEMENT_GRADIENT_BUILDER(GetReduceMeanGradient) {
std::vector<int64_t> axes_values(data_shape.size());
if (attributes.find("axes") != attributes.end()) {
axes_values = RetrieveValues<int64_t>(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<int64_t> 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<Dimension> data_shape = GetShape(I(0));
std::vector<NodeDef> result;
auto attributes = SrcNodeAttributes();
std::vector<int64_t> axes_values(data_shape.size());
if (attributes.find("axes") != attributes.end()) {
axes_values = RetrieveValues<int64_t>(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<bool>(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<int64_t> 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.");

View file

@ -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)

View file

@ -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);

View file

@ -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<int64_t>{-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<int64_t>{-2, -1}),
MakeAttribute("keepdims", int64_t(0))});
EXPECT_IS_TINY(max_error);
}
}
TEST(GradientCheckerTest, ReduceSumGrad) {
float max_error;
GradientChecker<float, float, float> 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<int64_t>{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<int64_t>{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<int64_t>{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<int64_t>{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<int64_t>{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<int64_t>{-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<int64_t>{-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<int>(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();
}