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 the BroadcastGradientArgs op (#4511)

Browse source 
* Adding CPU implementation of BroadcastGradientArgs op

* Modify to take shape as input instead of tensor

* Cleanup

* Correct schema

* Corrected kernel, added tests, addressed review comments.

* Added exception,test for invalid broadcast,addresed review comments.

* Fix mac build error.
This commit is contained in:
ashbhandare 2020-07-20 23:59:10 -07:00 committed by GitHub
parent 0008e92b4e
commit ab4be8355f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 219 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
orttraining/orttraining/core/framework/gradient_graph_builder.h
View file

@ -39,7 +39,8 @@ static std::unordered_map<std::string, std::unordered_set<size_t>>
{"Scatter", {1}},
{"OneHot", {0, 1, 2}},
{"Where", {0}},
{"Range", {0, 1, 2}}};
{"Range", {0, 1, 2}},
{"BroadcastGradientArgs", {0, 1}}};
class GradientGraphBuilder {
public:

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

@ -950,7 +950,7 @@ Example 4:
updateOutputShape(ctx, 0, TensorShapeProto());
}
if(ctx.getNumOutputs() == 2) {
if (ctx.getNumOutputs() == 2) {
propagateElemTypeFromInputToOutput(ctx, 0, 1);
if (hasInputShape(ctx, 0)) {
propagateShapeFromInputToOutput(ctx, 0, 1);
@ -1167,6 +1167,23 @@ Example 4:
propagateShapeAndTypeFromFirstInput(ctx);
});
ONNX_CONTRIB_OPERATOR_SCHEMA(BroadcastGradientArgs)
.SetDomain(kMSDomain)
.SinceVersion(1)
.SetSupportLevel(OpSchema::SupportType::EXPERIMENTAL)
.SetDoc(
"Returns the reduction axes for computing gradients of s0 op s1 with broadcast."
"The ouput axes are deterministic from last to first. "
"Output is an empty vector when no reduction is necessary for the corresponding input.")
.Input(0, "a_shape", "The 1st input shape as Tensor.", "T")
.Input(1, "b_shape", "The 2nd input shape as Tensor.", "T")
.Output(0, "a_axes", "The reduction axes for 1st input, last to first.", "T")
.Output(1, "b_axes", "The reduction axes for 2nd input, last to first.", "T")
.TypeConstraint(
"T",
{"tensor(int64)"},
"Constrain input and output types to 64-bit integer.");
ONNX_CONTRIB_OPERATOR_SCHEMA(GistBinarizeEncoder)
.SetDomain(kMSDomain)
.SinceVersion(1)

97
orttraining/orttraining/test/training_ops/cpu/nn/broadcast_grad_args_test.cc Normal file
View file

@ -0,0 +1,97 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <algorithm>
#include <memory>
#include <numeric>
#include <random>
#include "gtest/gtest.h"
#include "test/common/tensor_op_test_utils.h"
#include "test/providers/provider_test_utils.h"
#include "test/util/include/default_providers.h"
namespace onnxruntime {
namespace contrib {
namespace test {
using namespace onnxruntime::test;
namespace {
constexpr auto k_opset_version = 1;
void RunBroadcastGradientArgsTest(const char* op,
const std::vector<int64_t>& A_shape_tensor,
const std::vector<int64_t>& B_shape_tensor,
const std::vector<int64_t>& A_axes_expected,
const std::vector<int64_t>& B_axes_expected,
bool fail = false) {
OpTester t{op, k_opset_version, kMSDomain};
t.AddInput("a_shape", {static_cast<int64_t>(A_shape_tensor.size())}, A_shape_tensor);
t.AddInput("b_shape", {static_cast<int64_t>(B_shape_tensor.size())}, B_shape_tensor);
t.AddOutput<int64_t>("a_axes", {static_cast<int64_t>(A_axes_expected.size())}, A_axes_expected);
t.AddOutput<int64_t>("b_axes", {static_cast<int64_t>(B_axes_expected.size())}, B_axes_expected);
std::vector<std::unique_ptr<IExecutionProvider>> execution_providers;
execution_providers.push_back(DefaultCpuExecutionProvider());
if (fail)
t.Run(OpTester::ExpectResult::kExpectFailure, "", {}, nullptr, &execution_providers, ExecutionMode::ORT_SEQUENTIAL);
else
t.Run(OpTester::ExpectResult::kExpectSuccess, "", {}, nullptr, &execution_providers, ExecutionMode::ORT_SEQUENTIAL);
}
} // namespace
// BroadcastGradientArgs
TEST(BroadcastGradientArgsTest, Basic) {
RunBroadcastGradientArgsTest("BroadcastGradientArgs", {2, 16, 1024, 1024}, {1, 1, 1024, 1024},
{}, {1, 0});
}
TEST(BroadcastGradientArgsTest, Basic_both_valid_op) {
RunBroadcastGradientArgsTest("BroadcastGradientArgs", {2, 16, 1, 1024}, {1, 1, 1024, 1024},
{2}, {1, 0});
}
TEST(BroadcastGradientArgsTest, Basic_no_bcast) {
RunBroadcastGradientArgsTest("BroadcastGradientArgs", {2, 3, 4, 5}, {2, 3, 4, 5},
{}, {});
}
TEST(BroadcastGradientArgsTest, Basic_B_scalar) {
RunBroadcastGradientArgsTest("BroadcastGradientArgs", {2, 3, 4, 5}, {},
{}, {3, 2, 1, 0});
}
TEST(BroadcastGradientArgsTest, Basic_B_vector) {
RunBroadcastGradientArgsTest("BroadcastGradientArgs", {2, 3, 4, 5}, {5},
{}, {2, 1, 0});
}
TEST(BroadcastGradientArgsTest, Basic_A_bcast_different_size) {
RunBroadcastGradientArgsTest("BroadcastGradientArgs", {4, 5}, {2, 3, 4, 5},
{1, 0}, {});
}
TEST(BroadcastGradientArgsTest, Basic_both_bcast_different_size) {
RunBroadcastGradientArgsTest("BroadcastGradientArgs", {1, 4, 5}, {2, 3, 1, 1},
{1, 0}, {3, 2});
}
TEST(BroadcastGradientArgsTest, Basic_both_bcast_different_size_2) {
RunBroadcastGradientArgsTest("BroadcastGradientArgs", {3, 4, 5}, {2, 1, 1, 1},
{0}, {3, 2, 1});
}
TEST(BroadcastGradientArgsTest, Basic_invalid_broadcast) {
RunBroadcastGradientArgsTest("BroadcastGradientArgs", {3, 4, 5}, {2, 1, 6, 1},
{}, {}, true /*fail*/);
}
} // namespace test
} // namespace contrib
} // namespace onnxruntime

2
orttraining/orttraining/training_ops/cpu/cpu_training_kernels.cc
View file

@ -16,6 +16,7 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, AdamO
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, InPlaceAccumulator);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, ZeroGradient);
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Group);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int64_t, BroadcastGradientArgs);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ReduceSumTraining);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double, ReduceSumTraining);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int32_t, ReduceSumTraining);
@ -79,6 +80,7 @@ Status RegisterCpuTrainingKernels(KernelRegistry& kernel_registry) {
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, InPlaceAccumulator)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, ZeroGradient)>,
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, Group)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int64_t, BroadcastGradientArgs)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, float, ReduceSumTraining)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, double, ReduceSumTraining)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kMSDomain, 1, int32_t, ReduceSumTraining)>,

81
orttraining/orttraining/training_ops/cpu/nn/broadcast_grad_args.cc Normal file
View file

@ -0,0 +1,81 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "orttraining/training_ops/cpu/nn/broadcast_grad_args.h"
namespace onnxruntime {
namespace contrib {
#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_TYPED_KERNEL_EX( \
BroadcastGradientArgs, \
kMSDomain, \
1, \
T, \
kCpuExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
BroadcastGradientArgs<T>);
REGISTER_KERNEL_TYPED(int64_t)
template <typename T>
Status BroadcastGradientArgs<T>::Compute(OpKernelContext* context) const {
const Tensor* a_shape = context->Input<Tensor>(0);
const Tensor* b_shape = context->Input<Tensor>(1);
const T* A_dims = a_shape->template Data<T>();
const T* B_dims = b_shape->template Data<T>();
const T a_size = a_shape->Shape().Size();
const T b_size = b_shape->Shape().Size();
T ndim = std::max(a_size, b_size);
std::vector<T> a_axes, b_axes;
T i = a_size - 1;
T j = b_size - 1;
T k = ndim - 1;
for (; i >= 0 && j >= 0; --k) {
auto A_dim = A_dims[i],
B_dim = B_dims[j];
if (A_dim != B_dim) {
if (A_dim == 1) {
a_axes.push_back(k);
} else if (B_dim == 1) {
b_axes.push_back(k);
} else {
TensorShape a(A_dims, a_size);
TensorShape b(B_dims, b_size);
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL,
"Broadcast is not possible between inputs of shapes: ",
a, " and ", b);
}
}
--i;
--j;
}
if (i < 0) {
for (; k >= 0; --k) {
a_axes.push_back(k);
}
} else {
for (; k >= 0; --k) {
b_axes.push_back(k);
}
}
Tensor* A_axes = context->Output(0, {static_cast<T>(a_axes.size())});
T* A_axes_data = A_axes->template MutableData<T>();
std::copy(a_axes.begin(), a_axes.end(), A_axes_data);
Tensor* B_axes = context->Output(1, {static_cast<T>(b_axes.size())});
T* B_axes_data = B_axes->template MutableData<T>();
std::copy(b_axes.begin(), b_axes.end(), B_axes_data);
return Status::OK();
}
} // namespace contrib
} // namespace onnxruntime

19
orttraining/orttraining/training_ops/cpu/nn/broadcast_grad_args.h Normal file
View file

@ -0,0 +1,19 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include "core/framework/op_kernel.h"
namespace onnxruntime {
namespace contrib {
template <typename T>
class BroadcastGradientArgs final : public OpKernel {
public:
BroadcastGradientArgs(const OpKernelInfo& info) : OpKernel{info} {
}
Status Compute(OpKernelContext* context) const override;
};
} // namespace contrib
} // namespace onnxruntime