Complete GetGlobalAveragePoolGradient (#7514)

* Improve GetGlobalAveragePoolGradient

Co-authored-by: Sherlock Huang <bahuang@OrtTrainingDev3.af05slrtruoetgaxwwjv5nsq5e.px.internal.cloudapp.net>
This commit is contained in:
Sherlock 2021-04-30 18:04:01 -07:00 committed by GitHub
parent 9c1900866a
commit 668a65f1a7
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 45 additions and 29 deletions

View file

@ -12,6 +12,7 @@
#include "core/framework/tensorprotoutils.h"
#include "core/providers/common.h"
#include "core/common/safeint.h"
#include "orttraining/core/framework/distributed_run_context.h"
#include "orttraining/core/graph/gradient_builder_registry.h"
#include "orttraining/core/graph/graph_augmenter.h"
@ -1190,37 +1191,48 @@ IMPLEMENT_GRADIENT_BUILDER(GetSoftmaxCrossEntropyLossGradient) {
}
IMPLEMENT_GRADIENT_BUILDER(GetGlobalAveragePoolGradient) {
const ArgDef X = I(0);
const ArgDef X = I(0), Y = O(0), dX = GI(0), dY = GO(0);
// TODO: ONNX supports unknown shape for the input feed, e.g. [1, 3, -1, 28],
// thus the shape of input might be missing at graph construction time.
// However, in practice, we haven't seen a single model with unknown input shape.
// We need to get the shape at runtime if this case need to be supported.
// One way to do it is: scale = Size_Op(X, from=2); scaled_dY = Mul_Op(dY, scale)
const auto& x_dims = X.type_proto->tensor_type().shape().dim();
ORT_ENFORCE(x_dims.size() >= 3, "Input dimension cannot be less than 3.");
int64_t scale = 1;
for (auto dim = x_dims.begin() + 2; dim < x_dims.end(); dim++) {
if (dim->has_dim_value()) {
scale *= dim->dim_value();
} else {
ORT_ENFORCE(false, "Dimension missing");
bool has_concrete_shape = true;
SafeInt<int64_t> scale = 1;
std::vector<Dimension> x_dims;
if (GetShape(X, x_dims).IsOK()) {
ORT_ENFORCE(x_dims.size() >= 3, "Input dimension cannot be less than 3.");
for (auto dim = x_dims.begin() + 2; dim < x_dims.end(); dim++) {
if (dim->has_dim_value()) {
scale *= dim->dim_value();
} else {
has_concrete_shape = false;
break;
}
}
} else {
has_concrete_shape = false;
}
NodeDef scale_node = ConstantScalarNode(1.0f / static_cast<float>(scale), Name("Scale"), IElemType(0));
ArgDef SCALE = scale_node.output_args[0];
return std::vector<NodeDef>{
scale_node,
NodeDef("Mul",
{GO(0), SCALE},
{IA("scaled_dY")}),
NodeDef("Shape",
{X},
{IA("x_shape")}),
NodeDef("Expand",
{IA("scaled_dY"), IA("x_shape")},
{GI(0)})};
std::vector<NodeDef> result;
ArgDef scale_argdef;
if (has_concrete_shape) {
NodeDef scale_node = ConstantScalarNode(static_cast<float>(scale), Name("Scale"), IElemType(0));
result.push_back(scale_node);
scale_argdef = scale_node.output_args[0];
} else {
result.push_back(NodeDef("Size", {X}, {IA("X_Size")}));
result.push_back(NodeDef("Size", {Y}, {IA("Y_Size")}));
scale_argdef = IA("Scale");
result.push_back(NodeDef("Div", {IA("X_Size"), IA("Y_Size")}, {scale_argdef}));
}
result.push_back(NodeDef(OpDef{"Scale", kMSDomain, 1},
{dY, scale_argdef},
{IA("scaled_dY")},
{MakeAttribute("scale_down", int64_t(1))}));
result.push_back(NodeDef("Shape", {X}, {IA("x_shape")}));
result.push_back(NodeDef("Expand", {IA("scaled_dY"), IA("x_shape")}, {dX}));
return result;
}
IMPLEMENT_GRADIENT_BUILDER(GetGeluGradient) {

View file

@ -796,13 +796,17 @@ TEST(GradientCheckerTest, GlobalAveragePoolGrad) {
//globalaveragepool
{
gradient_checker.ComputeGradientError(op_def, {{2, 3, 5, 5}}, {{2, 3, 1, 1}}, &max_error);
gradient_checker.ComputeGradientError(op_def, {{2, 3, 5, 5}}, {{2, 3, 1, 1}}, &max_error, {},
/*check_not_have_gradient*/ true,
/*check_not_have_shape_inferencing*/ true);
EXPECT_IS_TINIER_THAN(max_error, error_tolerance);
}
//globalaveragepool_precomputed
{
gradient_checker.ComputeGradientError(op_def, {{2, 1, 3, 3}}, {{2, 1, 1, 1}}, &max_error);
gradient_checker.ComputeGradientError(op_def, {{2, 1, 3, 3}}, {{2, 1, 1, 1}}, &max_error, {},
/*check_not_have_gradient*/ true,
/*check_not_have_shape_inferencing*/ true);
EXPECT_IS_TINIER_THAN(max_error, error_tolerance);
}
}