mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-07 17:15:29 +00:00
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
This commit is contained in:
harshithapv
GitHub
parent
ce9cd6ad9a
commit
a5d3a52d1a
4 changed files with 140 additions and 19 deletions
103
orttraining/orttraining/core/graph/gradient_builder.cc
Normal file → Executable file
103
orttraining/orttraining/core/graph/gradient_builder.cc
Normal file → Executable file
|
|
@ -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>{
|
||||
NodeDef(OpDef{"Squeeze", kOnnxDomain, 13},
|
||||
{GO(0), I(1)},
|
||||
|
|
@ -790,7 +790,7 @@ IMPLEMENT_GRADIENT_BUILDER(GetReluGradient) {
|
|||
IMPLEMENT_GRADIENT_BUILDER(GetSqueezeGradient) {
|
||||
std::vector<NodeDef> result;
|
||||
size_t numInputs = GetSrcNodeInputSize();
|
||||
if (SrcNodeOpsetVersion() < 13) { //axes attribute
|
||||
if (SrcNodeOpsetVersion() < 13) { //axes attribute
|
||||
auto attributes = SrcNodeAttributes();
|
||||
std::vector<int64_t> 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>{
|
||||
NodeDef(OpDef{"FastGeluGrad", kMSDomain, 1},
|
||||
|
|
@ -1484,8 +1483,74 @@ IMPLEMENT_GRADIENT_BUILDER(GetClipGradient) {
|
|||
IMPLEMENT_GRADIENT_BUILDER(GetAbsGradient) {
|
||||
return std::vector<NodeDef>{
|
||||
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<NodeDef> result = {};
|
||||
|
||||
result.push_back(NodeDef("Shape", {I(0)}, {IA("orig_shape")}));
|
||||
std::vector<int64_t> 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<int64_t> 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<Dimension> 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<int64_t> 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
|
||||
|
|
|
|||
1
orttraining/orttraining/core/graph/gradient_builder.h
Normal file → Executable file
1
orttraining/orttraining/core/graph/gradient_builder.h
Normal file → Executable file
|
|
@ -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
|
||||
|
|
|
|||
1
orttraining/orttraining/core/graph/gradient_builder_registry.cc
Normal file → Executable file
1
orttraining/orttraining/core/graph/gradient_builder_registry.cc
Normal file → Executable file
|
|
@ -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
|
||||
|
|
|
|||
54
orttraining/orttraining/test/gradient/gradient_ops_test.cc
Normal file → Executable file
54
orttraining/orttraining/test/gradient/gradient_ops_test.cc
Normal file → Executable file
|
|
@ -2437,6 +2437,60 @@ TEST(GradientCheckerTest, ClipGrad) {
|
|||
}
|
||||
}
|
||||
|
||||
TEST(GradientCheckerTest, TileGrad) {
|
||||
float max_error;
|
||||
GradientChecker<float, float, float> gradient_checker;
|
||||
OpDef op_def{"Tile", kOnnxDomain, 11};
|
||||
|
||||
// 2D input
|
||||
{
|
||||
TensorInfo x_info({2, 4}, true);
|
||||
TensorInfo repeat_info({2}, false, nullptr, DataTypeImpl::GetTensorType<int64_t>());
|
||||
std::vector<std::vector<float>> 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<int64_t>());
|
||||
std::vector<std::vector<float>> 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<int64_t>());
|
||||
std::vector<std::vector<float>> 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<int64_t>());
|
||||
std::vector<std::vector<float>> 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
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue