saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-11 17:48:34 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add fallback function implementation for DivGrad (#5518)

Browse source 
* Add fallback function implementation for DivGrad.

* Add shape inference for DivGrad.

* Add missing argument.

Co-authored-by: Derek Murray <demurra@microsoft.com>
This commit is contained in:
Derek Murray 2020-10-19 10:47:47 -07:00 committed by GitHub
parent a355281b99
commit 0b59004666
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 36 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

37
orttraining/orttraining/core/graph/training_op_defs.cc
View file

@ -501,7 +501,42 @@ void RegisterTrainingOpSchemas() {
.TypeConstraint(
"T",
{"tensor(float16)", "tensor(float)", "tensor(double)", "tensor(bfloat16)"},
"Constrain input and output types to numeric tensors.");
"Constrain input and output types to numeric tensors.")
.FunctionBody([]() {
auto nodes = ONNX_NAMESPACE::FunctionBodyHelper::BuildNodes(
{// nodes: {outputs, op, inputs, attributes}
// Get input shapes and dynamic reduction axes.
{{"shape_A"}, "Shape", {"A"}},
{{"shape_B"}, "Shape", {"B"}},
{{"axes_A", "axes_B"}, "BroadcastGradientArgs", {"shape_A", "shape_B"}},
// dA = reshape(reduce_sum(dY / B, axes_A), shape_A)
{{"dY_over_B"}, "Div", {"dY", "B"}},
{{"reduce_dA"}, "ReduceSumTraining", {"dY_over_B", "axes_A"},
{ONNX_NAMESPACE::MakeAttribute("noop_with_empty_axes", int64_t(1))}},
{{"dA"}, "Reshape", {"reduce_dA", "shape_A"}},
// dB = reshape(reduce_sum(dY * -A / (B * B)), axes_B), shape_B)
{{"B_squared"}, "Mul", {"B", "B"}},
{{"minus_A"}, "Neg", {"A"}},
{{"minus_A_over_B_squared"}, "Div", {"minus_A", "B_squared"}},
{{"pre_reduce_dB"}, "Mul", {"dY", "minus_A_over_B_squared"}},
{{"reduce_dB"}, "ReduceSumTraining", {"pre_reduce_dB", "axes_B"},
{ONNX_NAMESPACE::MakeAttribute("noop_with_empty_axes", int64_t(1))}},
{{"dB"}, "Reshape", {"reduce_dB", "shape_B"}}});
for (size_t contrib_node_index : {2, 4, 10}) {
nodes[contrib_node_index].set_domain(kMSDomain);
}
return nodes;
}())
.TypeAndShapeInferenceFunction([](ONNX_NAMESPACE::InferenceContext& ctx) {
for (size_t i = 0; i < ctx.getNumOutputs(); ++i) {
propagateElemTypeFromTensorInputToOutput(ctx, 0, i);
propagateShapeFromInputToOutput(ctx, i + 1, i);
}
});
//TODO: Move this to the right location. Its only here for quick experimentation.
//TODO: Use the mutli weight / grad version.

2
orttraining/orttraining/test/gradient/gradient_ops_test.cc
View file

@ -294,12 +294,10 @@ TEST(GradientCheckerTest, DISABLED_MulGrad) {
TestBroadcastableBinaryOpGrad("Mul");
}
#ifdef USE_CUDA
TEST(GradientCheckerTest, DivGrad) {
std::function<float(float)> transformer = [](float x) { return x > 0 ? x + 0.2f : x - 0.2f; };
TestBroadcastableBinaryOpGrad("Div", &transformer);
}
#endif
// TODO: Powgrad Test doesn't cover exponent
TEST(GradientCheckerTest, PowGrad) {