mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-27 20:02:15 +00:00
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:
parent
e7297e6c9d
commit
03996c7c08
6 changed files with 208 additions and 28 deletions
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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];
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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.");
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in a new issue