From 6d8e81cb08ead9115bcf8d8936c56a60141cf51b Mon Sep 17 00:00:00 2001 From: ashbhandare Date: Wed, 4 Nov 2020 20:00:27 -0800 Subject: [PATCH] Update Squeeze, Unsqueeze, Split and ReduceSum kernel for Opset13 (#5691) * Split change * ReduceSum and Split change * Other op changes, Grad builder, tests, registering required opset 13 ops * Rebase fixes * Fix tests, add some more * Review changes, rebase * Fix windows build * Disable new tests for TesnorRT EP * Disable unsupported for OpenVINO Co-authored-by: Aishwarya --- .../providers/cpu/cpu_execution_provider.cc | 56 ++-- .../providers/cpu/reduction/reduction_ops.cc | 11 +- .../core/providers/cpu/tensor/concat.cc | 1 + .../core/providers/cpu/tensor/shape_op.cc | 1 - .../core/providers/cpu/tensor/split.cc | 42 ++- onnxruntime/core/providers/cpu/tensor/split.h | 15 +- .../core/providers/cpu/tensor/squeeze.cc | 12 +- .../core/providers/cpu/tensor/squeeze.h | 42 ++- .../core/providers/cpu/tensor/unsqueeze.cc | 36 ++- .../core/providers/cpu/tensor/unsqueeze.h | 5 +- .../providers/cuda/reduction/reduction_ops.cc | 83 ++++- .../core/providers/cuda/tensor/split.cc | 18 +- .../core/providers/cuda/tensor/squeeze.cc | 21 +- .../core/providers/cuda/tensor/unsqueeze.cc | 7 +- .../cpu/reduction/reduction_ops_test.cc | 55 ++++ .../providers/cpu/tensor/split_op_test.cc | 70 +++- .../providers/cpu/tensor/squeeze_op_test.cc | 35 ++ .../providers/cpu/tensor/unsqueeze_op_test.cc | 98 ++++-- .../core/graph/gradient_builder.cc | 78 +++-- .../core/graph/gradient_builder_base.h | 4 + .../test/gradient/gradient_ops_test.cc | 305 ++++++++++-------- 21 files changed, 729 insertions(+), 266 deletions(-) diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index 8bc88f0710..ecafbc0a70 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc @@ -326,10 +326,10 @@ class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOn class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, float, ReduceProd); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, int32_t, ReduceProd); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, int64_t, ReduceProd); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, float, ReduceSum); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, double, ReduceSum); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, int32_t, ReduceSum); -class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, int64_t, ReduceSum); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, float, ReduceSum); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, double, ReduceSum); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, int32_t, ReduceSum); +class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, int64_t, ReduceSum); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, float, ReduceSumSquare); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, double, ReduceSumSquare); class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, int32_t, ReduceSumSquare); @@ -346,9 +346,9 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Co class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, Concat); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, Gather); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, Slice); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Split); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Squeeze); -class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Unsqueeze); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, Split); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, Squeeze); +class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, Unsqueeze); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, Det); class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, 12, ScatterElements); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 11, NonMaxSuppression); @@ -529,6 +529,9 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Po class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, DepthToSpace); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, SpaceToDepth); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Slice); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Split); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Unsqueeze); +class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Squeeze); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Transpose); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Tile); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, Gather); @@ -571,6 +574,10 @@ class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, float, ReduceSumSquare); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, double, ReduceSumSquare); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int32_t, ReduceSumSquare); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, float, ReduceSum); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, double, ReduceSum); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int32_t, ReduceSum); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int64_t, ReduceSum); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, float, Resize); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, int32_t, Resize); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 13, uint8_t, Resize); @@ -1048,9 +1055,9 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, - BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo, @@ -1455,6 +1465,14 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) { ReduceSumSquare)>, BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, + BuildKernelCreateInfo, BuildKernelCreateInfo, BuildKernelCreateInfo CPUExecutionProvider::GetKernelRegistry() const std::unique_ptr CPUExecutionProvider::GetDataTransfer() const { return onnxruntime::make_unique(); } -} // namespace onnxruntime +} // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc b/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc index db095a4498..14dc9151e0 100644 --- a/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cpu/reduction/reduction_ops.cc @@ -166,11 +166,14 @@ REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceProd, 13); REGISTER_UNARY_ELEMENTWISE_KERNEL_INT64_ONLY(ReduceProd, 13); REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ReduceSum, 1, 10); -REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceSum, 11); REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL_INT64_ONLY(ReduceSum, 1, 10); -REGISTER_UNARY_ELEMENTWISE_KERNEL_INT64_ONLY(ReduceSum, 11); REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL_DOUBLE_ONLY(ReduceSum, 1, 10); -REGISTER_UNARY_ELEMENTWISE_KERNEL_DOUBLE_ONLY(ReduceSum, 11); +REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ReduceSum, 11, 12); +REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL_INT64_ONLY(ReduceSum, 11, 12); +REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL_DOUBLE_ONLY(ReduceSum, 11, 12); +REGISTER_UNARY_ELEMENTWISE_KERNEL(ReduceSum, 13); +REGISTER_UNARY_ELEMENTWISE_KERNEL_INT64_ONLY(ReduceSum, 13); +REGISTER_UNARY_ELEMENTWISE_KERNEL_DOUBLE_ONLY(ReduceSum, 13); REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL(ReduceSumSquare, 1, 10); REGISTER_UNARY_ELEMENTWISE_VERSIONED_KERNEL_DOUBLE_ONLY(ReduceSumSquare, 1, 10); @@ -560,7 +563,7 @@ Status ReduceProd::Compute(OpKernelContext* ctx) const { template Status ReduceSum::Compute(OpKernelContext* ctx) const { ResultsNoTransposePrepareForReduce last_results; - CommonReduce>(ctx, axes_, keepdims_, last_results); + CommonReduce>(ctx, axes_, keepdims_, last_results, noop_with_empty_axes_); return Status::OK(); } diff --git a/onnxruntime/core/providers/cpu/tensor/concat.cc b/onnxruntime/core/providers/cpu/tensor/concat.cc index af058b1fbd..8fd49ffdf7 100644 --- a/onnxruntime/core/providers/cpu/tensor/concat.cc +++ b/onnxruntime/core/providers/cpu/tensor/concat.cc @@ -22,6 +22,7 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllTensorTypes()), Concat); +// Opset 13 . ONNX_CPU_OPERATOR_KERNEL( Concat, 13, diff --git a/onnxruntime/core/providers/cpu/tensor/shape_op.cc b/onnxruntime/core/providers/cpu/tensor/shape_op.cc index 76d491599d..a089ccdfe9 100644 --- a/onnxruntime/core/providers/cpu/tensor/shape_op.cc +++ b/onnxruntime/core/providers/cpu/tensor/shape_op.cc @@ -17,5 +17,4 @@ ONNX_CPU_OPERATOR_KERNEL( 13, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()).TypeConstraint("T1", DataTypeImpl::GetTensorType()), Shape); - } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/tensor/split.cc b/onnxruntime/core/providers/cpu/tensor/split.cc index 2af92504ec..0d3b7ca54c 100644 --- a/onnxruntime/core/providers/cpu/tensor/split.cc +++ b/onnxruntime/core/providers/cpu/tensor/split.cc @@ -23,9 +23,22 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( Split); // Opset 11 starts to support Neg Axis. -ONNX_CPU_OPERATOR_KERNEL( +ONNX_CPU_OPERATOR_VERSIONED_KERNEL( Split, 11, + 12, + KernelDefBuilder().TypeConstraint("T", + std::vector{ + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType(), + DataTypeImpl::GetTensorType()}), + Split); + +// Opset 13 starts to supports 'split' as optional input. +ONNX_CPU_OPERATOR_KERNEL( + Split, + 13, KernelDefBuilder().TypeConstraint("T", std::vector{ DataTypeImpl::GetTensorType(), @@ -48,7 +61,7 @@ Status SplitBase::PrepareForCompute(const TensorShape& input_shape, int num_outp ? 1 // we multiply by this value so must be 1 not 0 : gsl::narrow(input_shape.SizeFromDimension(axis + 1)); - if (split_sizes_.empty()) { + if (split_sizes.empty()) { // equal split based on number of outputs if (split_dim_size % static_cast(num_outputs) != 0) { return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Input cannot be split evenly on selected axis. Input shape=", input_shape, @@ -58,17 +71,18 @@ Status SplitBase::PrepareForCompute(const TensorShape& input_shape, int num_outp // populate split_sizes with the same size for each output split_sizes = std::vector(static_cast(num_outputs), split_dim_size / num_outputs); } else { - if (split_sizes_.size() != static_cast(num_outputs) || split_size_sum_ != split_dim_size) + int64_t split_size_sum = split_size_sum_; + if (split_size_sum == -1){ + split_size_sum = std::accumulate(split_sizes.cbegin(), split_sizes.cend(), 0LL); + } + if (split_sizes.size() != static_cast(num_outputs) || split_size_sum != split_dim_size) return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Cannot split using values in 'split' attribute. Axis=", axis_, " Input shape=", input_shape, " NumOutputs=", num_outputs, - " Num entries in 'split' (must equal number of outputs) was ", split_sizes_.size(), - " Sum of sizes in 'split' (must equal size of selected axis) was ", split_size_sum_); - - split_sizes = split_sizes_; + " Num entries in 'split' (must equal number of outputs) was ", split_sizes.size(), + " Sum of sizes in 'split' (must equal size of selected axis) was ", split_size_sum); } - return Status::OK(); } @@ -112,6 +126,18 @@ Status Split::ComputeImpl(OpKernelContext& context, const Tensor& input) const { int after_dims_excluding_split = 0; std::vector split_sizes; + size_t num_inputs = context.InputCount(); + if (num_inputs == 2) { + //override the attribute value with the input value for split_split + const Tensor* split_tensor = context.Input(1); + ORT_ENFORCE(split_tensor->Shape().NumDimensions() == 1, "An split tensor must be a vector tensor."); + auto nDims = static_cast(split_tensor->Shape()[0]); + const auto* data = split_tensor->template Data(); + split_sizes.assign(data, data + nDims); + } + else{ + split_sizes.assign(split_sizes_.begin(), split_sizes_.end()); + } ORT_RETURN_IF_ERROR(PrepareForCompute(input_shape, num_outputs, axis, diff --git a/onnxruntime/core/providers/cpu/tensor/split.h b/onnxruntime/core/providers/cpu/tensor/split.h index d2e51230a7..ff6dcacdb3 100644 --- a/onnxruntime/core/providers/cpu/tensor/split.h +++ b/onnxruntime/core/providers/cpu/tensor/split.h @@ -15,11 +15,14 @@ class SplitBase { SplitBase(const OpKernelInfo& info) { axis_ = info.GetAttrOrDefault("axis", 0); - // optional - if (info.GetAttrs("split", split_sizes_).IsOK()) { - split_size_sum_ = std::accumulate(split_sizes_.cbegin(), split_sizes_.cend(), 0LL); - ORT_ENFORCE(std::all_of(split_sizes_.cbegin(), split_sizes_.cend(), [](int64_t value) { return value >= 0; }), - "Invalid value in 'split' attribute. All values must be > 0"); + size_t numInputs = info.GetInputCount(); + if(numInputs == 1){ + // optional + if (info.GetAttrs("split", split_sizes_).IsOK()) { + split_size_sum_ = std::accumulate(split_sizes_.cbegin(), split_sizes_.cend(), 0LL); + ORT_ENFORCE(std::all_of(split_sizes_.cbegin(), split_sizes_.cend(), [](int64_t value) { return value >= 0; }), + "Invalid value in 'split' attribute. All values must be > 0"); + } } } @@ -32,7 +35,7 @@ class SplitBase { int64_t axis_; std::vector split_sizes_; - int64_t split_size_sum_ = 0; + int64_t split_size_sum_ = -1; }; class Split final : public OpKernel, public SplitBase { diff --git a/onnxruntime/core/providers/cpu/tensor/squeeze.cc b/onnxruntime/core/providers/cpu/tensor/squeeze.cc index defcb83873..ef4e500e8b 100644 --- a/onnxruntime/core/providers/cpu/tensor/squeeze.cc +++ b/onnxruntime/core/providers/cpu/tensor/squeeze.cc @@ -15,9 +15,19 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( Squeeze); // Opset 11 starts to support Neg Axis. -ONNX_CPU_OPERATOR_KERNEL( +ONNX_CPU_OPERATOR_VERSIONED_KERNEL( Squeeze, 11, + 12, + KernelDefBuilder() + .TypeConstraint("T", DataTypeImpl::AllTensorTypes()) + .Alias(0, 0), + Squeeze); + +// axes is input instead of attribute +ONNX_CPU_OPERATOR_KERNEL( + Squeeze, + 13, KernelDefBuilder() .TypeConstraint("T", DataTypeImpl::AllTensorTypes()) .Alias(0, 0), diff --git a/onnxruntime/core/providers/cpu/tensor/squeeze.h b/onnxruntime/core/providers/cpu/tensor/squeeze.h index 7a3d82f436..c3d2544c40 100644 --- a/onnxruntime/core/providers/cpu/tensor/squeeze.h +++ b/onnxruntime/core/providers/cpu/tensor/squeeze.h @@ -14,27 +14,30 @@ class SqueezeBase { protected: explicit SqueezeBase(const OpKernelInfo& info) { std::vector axes; - // Parse attribute 'axes' - Status status = info.GetAttrs("axes", axes); + size_t numInputs = info.GetInputCount(); + if (numInputs == 1) { + // Parse attribute 'axes' + Status status = info.GetAttrs("axes", axes); - // Handle out of order and repeating dims when 'axes' exists. - if (status.IsOK()) { - std::sort(axes.begin(), axes.end()); - axes.erase(std::unique(axes.begin(), axes.end()), axes.end()); - axes_ = axes; + // Handle out of order and repeating dims when 'axes' exists. + if (status.IsOK()) { + std::sort(axes.begin(), axes.end()); + axes.erase(std::unique(axes.begin(), axes.end()), axes.end()); + axes_ = axes; + } } } static std::vector ComputeOutputShape( const TensorShape& input_shape, - const TensorShape& axes) { + const std::vector& axes) { size_t j = 0; std::vector output_shape; auto num_dimensions = input_shape.NumDimensions(); // Handle negtive axis, then resort and uniq. - std::vector axes_corrected(axes.NumDimensions()); - for (size_t i = 0; i < axes.NumDimensions(); i++) { + std::vector axes_corrected(axes.size()); + for (size_t i = 0; i < axes.size(); i++) { axes_corrected[i] = HandleNegativeAxis(axes[i], num_dimensions); } std::sort(axes_corrected.begin(), axes_corrected.end()); @@ -53,7 +56,7 @@ class SqueezeBase { return output_shape; } - TensorShape axes_; + std::vector axes_; }; class Squeeze final : public OpKernel, public SqueezeBase { @@ -63,7 +66,22 @@ class Squeeze final : public OpKernel, public SqueezeBase { Status Compute(OpKernelContext* context) const override { const auto* X = context->Input(0); const TensorShape& X_shape = X->Shape(); - std::vector output_shape = ComputeOutputShape(X_shape.GetDims(), axes_); + + std::vector axes; + size_t num_inputs = context->InputCount(); + if (num_inputs == 2) { //axes is an input + const Tensor* axes_tensor = context->Input(1); + ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); + ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, + "An axes tensor must be a vector tensor."); + auto nDims = static_cast(axes_tensor->Shape()[0]); + const auto* data = axes_tensor->template Data(); + axes.assign(data, data + nDims); + } else { + axes.assign(axes_.begin(), axes_.end()); + } + + std::vector output_shape = ComputeOutputShape(X_shape, axes); Tensor* Y = context->Output(0, TensorShape(output_shape)); diff --git a/onnxruntime/core/providers/cpu/tensor/unsqueeze.cc b/onnxruntime/core/providers/cpu/tensor/unsqueeze.cc index 39be162f0b..436f553eb8 100644 --- a/onnxruntime/core/providers/cpu/tensor/unsqueeze.cc +++ b/onnxruntime/core/providers/cpu/tensor/unsqueeze.cc @@ -18,9 +18,19 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL( .TypeConstraint("T", DataTypeImpl::AllTensorTypes()), Unsqueeze); -ONNX_CPU_OPERATOR_KERNEL( +ONNX_CPU_OPERATOR_VERSIONED_KERNEL( Unsqueeze, 11, + 12, + KernelDefBuilder() + .Alias(0, 0) + .TypeConstraint("T", DataTypeImpl::AllTensorTypes()), + Unsqueeze); + +// axes is input instead of attribute +ONNX_CPU_OPERATOR_KERNEL( + Unsqueeze, + 13, KernelDefBuilder() .Alias(0, 0) .TypeConstraint("T", DataTypeImpl::AllTensorTypes()), @@ -31,12 +41,26 @@ Status UnsqueezeBase::PrepareCompute(OpKernelContext* ctx, Prepare& p) const { ORT_ENFORCE(X != nullptr); auto& input_tensor = *X; - // New dimension count is the current dimensions + the number of entries in axes_ - // Initialize output_dims to 0 in each axis initially - std::vector output_dims(axes_.size() + input_tensor.Shape().NumDimensions(), 0); + std::vector axes; + size_t num_inputs = ctx->InputCount(); + if (num_inputs == 2) { //axes is an input + const Tensor* axes_tensor = ctx->Input(1); + ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); + ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, + "An axes tensor must be a vector tensor."); + auto nDims = static_cast(axes_tensor->Shape()[0]); + const auto* data = axes_tensor->template Data(); + axes.assign(data, data + nDims); + } else { + axes.assign(axes_.begin(), axes_.end()); + } - // Set all axes_ indices to 1 in output_dims and check for duplicates - for (int64_t axis : axes_) { + // New dimension count is the current dimensions + the number of entries in axes + // Initialize output_dims to 0 in each axis initially + std::vector output_dims(axes.size() + input_tensor.Shape().NumDimensions(), 0); + + // Set all axes indices to 1 in output_dims and check for duplicates + for (int64_t axis : axes) { // Valid axis range is [0, output_rank - 1] axis = HandleNegativeAxis(axis, output_dims.size()); if (axis < 0 || axis >= static_cast(output_dims.size())) diff --git a/onnxruntime/core/providers/cpu/tensor/unsqueeze.h b/onnxruntime/core/providers/cpu/tensor/unsqueeze.h index 3ab5d8d0f8..f5d18426c3 100644 --- a/onnxruntime/core/providers/cpu/tensor/unsqueeze.h +++ b/onnxruntime/core/providers/cpu/tensor/unsqueeze.h @@ -11,7 +11,10 @@ namespace onnxruntime { class UnsqueezeBase { protected: UnsqueezeBase(const OpKernelInfo& info) { - ORT_ENFORCE(info.GetAttrs("axes", axes_).IsOK(), "Missing/Invalid 'axes' attribute value"); + size_t numInputs = info.GetInputCount(); + if (numInputs == 1) { //axes must be a valid attribute + ORT_ENFORCE(info.GetAttrs("axes", axes_).IsOK(), "Missing/Invalid 'axes' attribute value"); + } } struct Prepare { diff --git a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc index 079a66b8f2..47fc6026a4 100644 --- a/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc +++ b/onnxruntime/core/providers/cuda/reduction/reduction_ops.cc @@ -95,6 +95,35 @@ namespace cuda { KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), \ name); +// Register with the latest version 13 +#define REGISTER_KERNEL_TYPED_13(name, T) \ + ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ + name, \ + kOnnxDomain, \ + 1, 10, \ + T, \ + kCudaExecutionProvider, \ + KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + name); \ + ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX( \ + name, \ + kOnnxDomain, \ + 11, 12, \ + T, \ + kCudaExecutionProvider, \ + KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + name); \ + ONNX_OPERATOR_TYPED_KERNEL_EX( \ + name, \ + kOnnxDomain, \ + 13, \ + T, \ + kCudaExecutionProvider, \ + KernelDefBuilder() \ + .InputMemoryType(1) \ + .TypeConstraint("T", DataTypeImpl::GetTensorType()), \ + name); + // TODO ReduceKernel::ReduceKernelShared() is still used by some other training classes though it's not used here - this should be refactored. template template @@ -618,11 +647,32 @@ template template Status ReduceKernel::ComputeImpl(OpKernelContext* ctx, cudnnReduceTensorOp_t cudnn_reduce_op) const { const Tensor* X = ctx->Input(0); + std::vector axes; + + size_t num_inputs = ctx->InputCount(); + if (num_inputs == 2) { + //override the attribute value with the input value for reduction_axes + const Tensor* axes_tensor = ctx->Input(1); + ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); + ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor."); + auto nDims = static_cast(axes_tensor->Shape()[0]); + const auto* data = axes_tensor->template Data(); + axes.assign(data, data + nDims); + } else { + axes.assign(axes_.begin(), axes_.end()); + } + + // empty axes and no-op + if (axes.empty() && noop_with_empty_axes_) { + auto* Y = ctx->Output(0, X->Shape()); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData(), X->template Data(), X->SizeInBytes(), cudaMemcpyDeviceToDevice)); + return Status::OK(); + } PrepareReduceMetadata prepare_reduce_metadata; ORT_RETURN_IF_ERROR(PrepareForReduce(X, keepdims_, - axes_, + axes, prepare_reduce_metadata)); Tensor* Y = ctx->Output(0, prepare_reduce_metadata.squeezed_output_dims); bool fast_reduction = fast_reduction_; @@ -642,12 +692,33 @@ Status ReduceKernel::ComputeImpl( typedef typename ToCudaType::MappedType CudaT; const Tensor* X = ctx->Input(0); + std::vector axes; + + size_t num_inputs = ctx->InputCount(); + if (num_inputs == 2) { + //override the attribute value with the input value for reduction_axes + const Tensor* axes_tensor = ctx->Input(1); + ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); + ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, "An axes tensor must be a vector tensor."); + auto nDims = static_cast(axes_tensor->Shape()[0]); + const auto* data = axes_tensor->template Data(); + axes.assign(data, data + nDims); + } else { + axes.assign(axes_.begin(), axes_.end()); + } + + // empty axes and no-op + if (axes.empty() && noop_with_empty_axes_) { + auto* Y = ctx->Output(0, X->Shape()); + CUDA_RETURN_IF_ERROR(cudaMemcpyAsync(Y->template MutableData(), X->template Data(), X->SizeInBytes(), cudaMemcpyDeviceToDevice)); + return Status::OK(); + } PrepareReduceMetadata prepare_reduce_metadata; ORT_RETURN_IF_ERROR(PrepareForReduce(X, keepdims_, - axes_, + axes, prepare_reduce_metadata)); Tensor* Y = ctx->Output(0, prepare_reduce_metadata.squeezed_output_dims); @@ -952,7 +1023,12 @@ REGISTER_KERNEL_TYPED_12(ReduceMin, int8_t) REGISTER_KERNEL_TYPED_12(ReduceMin, uint8_t) REGISTER_KERNEL_HFD(ReduceProd) -REGISTER_KERNEL_HFD(ReduceSum) + +REGISTER_KERNEL_TYPED_13(ReduceSum, MLFloat16) +REGISTER_KERNEL_TYPED_13(ReduceSum, float) +REGISTER_KERNEL_TYPED_13(ReduceSum, double) +REGISTER_KERNEL_TYPED_13(ReduceSum, int32_t) + REGISTER_KERNEL_HFD(ReduceLogSum) REGISTER_KERNEL_HFD(ReduceSumSquare) REGISTER_KERNEL_HFD(ReduceLogSumExp) @@ -965,7 +1041,6 @@ REGISTER_KERNEL_INT32(ReduceL2) REGISTER_KERNEL_INT32(ReduceMean) REGISTER_KERNEL_INT32(ReduceProd) -REGISTER_KERNEL_INT32(ReduceSum) } // namespace cuda } // namespace onnxruntime diff --git a/onnxruntime/core/providers/cuda/tensor/split.cc b/onnxruntime/core/providers/cuda/tensor/split.cc index f5919f22ae..539f8016a6 100644 --- a/onnxruntime/core/providers/cuda/tensor/split.cc +++ b/onnxruntime/core/providers/cuda/tensor/split.cc @@ -23,11 +23,14 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX(Split, KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), Split); +// explicitly supports 'split' as optional input ONNX_OPERATOR_KERNEL_EX(Split, kOnnxDomain, 13, kCudaExecutionProvider, - KernelDefBuilder().TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), + KernelDefBuilder() + .InputMemoryType(1) + .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), Split); Status Split::ComputeInternal(OpKernelContext* ctx) const { @@ -39,7 +42,18 @@ Status Split::ComputeInternal(OpKernelContext* ctx) const { int before_dims = 0; int block_size_including_axis_dim = 0; int block_size_inside_axis_dim = 0; - std::vector split_sizes; + std::vector split_sizes(num_outputs); + + size_t num_inputs = ctx->InputCount(); + if (num_inputs == 2) { + const Tensor* split_tensor = ctx->Input(1); + ORT_ENFORCE(split_tensor->Shape().NumDimensions() == 1, "An split tensor must be a vector tensor."); + auto nDims = static_cast(split_tensor->Shape()[0]); + const int64_t* data = split_tensor->template Data(); + split_sizes.assign(data, data + nDims); + } else { + split_sizes.assign(split_sizes_.begin(), split_sizes_.end()); + } ORT_RETURN_IF_ERROR(PrepareForCompute(input_shape, num_outputs, diff --git a/onnxruntime/core/providers/cuda/tensor/squeeze.cc b/onnxruntime/core/providers/cuda/tensor/squeeze.cc index f8819af6ce..b6cd7317ea 100644 --- a/onnxruntime/core/providers/cuda/tensor/squeeze.cc +++ b/onnxruntime/core/providers/cuda/tensor/squeeze.cc @@ -27,6 +27,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), Squeeze); +// axes is input instead of attribute ONNX_OPERATOR_KERNEL_EX( Squeeze, kOnnxDomain, @@ -34,13 +35,29 @@ ONNX_OPERATOR_KERNEL_EX( kCudaExecutionProvider, KernelDefBuilder() .Alias(0, 0) - .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), + .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()) + .InputMemoryType(1), Squeeze); Status Squeeze::ComputeInternal(OpKernelContext* ctx) const { const Tensor* X = ctx->Input(0); const TensorShape& X_shape = X->Shape(); - std::vector output_shape = ComputeOutputShape(X_shape, axes_); + + std::vector axes; + size_t num_inputs = ctx->InputCount(); + if (num_inputs == 2) { //axes is an input + const Tensor* axes_tensor = ctx->Input(1); + ORT_ENFORCE(axes_tensor != nullptr, "Axes input is null"); + ORT_ENFORCE(axes_tensor->Shape().NumDimensions() == 1, + "An axes tensor must be a vector tensor."); + auto nDims = static_cast(axes_tensor->Shape()[0]); + const auto* data = axes_tensor->template Data(); + axes.assign(data, data + nDims); + } else { + axes.assign(axes_.begin(), axes_.end()); + } + + std::vector output_shape = ComputeOutputShape(X_shape, axes); Tensor* Y = ctx->Output(0, TensorShape(output_shape)); diff --git a/onnxruntime/core/providers/cuda/tensor/unsqueeze.cc b/onnxruntime/core/providers/cuda/tensor/unsqueeze.cc index 152fadda06..c528322d0e 100644 --- a/onnxruntime/core/providers/cuda/tensor/unsqueeze.cc +++ b/onnxruntime/core/providers/cuda/tensor/unsqueeze.cc @@ -27,7 +27,7 @@ ONNX_OPERATOR_VERSIONED_KERNEL_EX( .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), Unsqueeze); -// support bfloat16 +// axes is input instead of attribute, support bfloat16 ONNX_OPERATOR_KERNEL_EX( Unsqueeze, kOnnxDomain, @@ -35,8 +35,9 @@ ONNX_OPERATOR_KERNEL_EX( kCudaExecutionProvider, KernelDefBuilder() .Alias(0, 0) - .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()), - Unsqueeze); + .TypeConstraint("T", DataTypeImpl::AllFixedSizeTensorTypes()) + .InputMemoryType(1), + Unsqueeze); Status Unsqueeze::ComputeInternal(OpKernelContext* ctx) const { Prepare p; diff --git a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc index 8c3ab7981d..da337c0db1 100644 --- a/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc +++ b/onnxruntime/test/providers/cpu/reduction/reduction_ops_test.cc @@ -1330,6 +1330,61 @@ TEST(ReductionOpTest, ReduceSum_keepdims) { test.Run(); } +TEST(ReductionOpTest, ReduceSum_int32_axes_input) { + OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)1); + test.AddInput("data", {3, 2, 2}, + {1, 2, + 3, 4, + + 5, 6, + 7, 8, + + 9, 10, + 11, 12}); + test.AddInput("axes", {2}, std::vector{0, 2}, true); + test.AddOutput("reduced", {1, 2, 1}, {33, 45}); + // TODO: TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); +} + +TEST(ReductionOpTest, ReduceSum_do_not_keepdims_axes_input_initializer) { + OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {1, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f}); + test.AddInput("axes", {1}, std::vector{1}, true); + test.AddOutput("reduced", {1, 2}, {4.0f, 6.0f}); + // TODO: TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); +} + +TEST(ReductionOpTest, ReduceSum_do_not_keepdims_axes_input_not_initializer) { + OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)0); + test.AddInput("data", {1, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f}); + test.AddInput("axes", {1}, std::vector{1}, false); + test.AddOutput("reduced", {1, 2}, {4.0f, 6.0f}); + // TODO: TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); +} + +TEST(ReductionOpTest, ReduceSum_noop_axes_input_initializer) { + OpTester test("ReduceSum", 13, onnxruntime::kOnnxDomain); + test.AddAttribute("keepdims", (int64_t)0); + test.AddAttribute("noop_with_empty_axes", (int64_t)1); + test.AddInput("data", {1, 2, 2}, + {1.0f, 2.0f, + 3.0f, 4.0f}); + test.AddInput("axes", {0}, {}, true); + test.AddOutput("reduced", {1, 2, 2}, {1.0f, 2.0f, 3.0f, 4.0f}); + // TODO: TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); +} + #if !(defined USE_TENSORRT) && !(defined USE_TVM) TEST(ReductionOpTest, ReduceSum0DTensor) { OpTester test("ReduceSum"); diff --git a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc index f9f8561144..cff6453d69 100644 --- a/onnxruntime/test/providers/cpu/tensor/split_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/split_op_test.cc @@ -17,15 +17,19 @@ using ExpectResult = OpTester::ExpectResult; template void RunTest(int64_t axis, const std::vector split_sizes, const ShapeAndData& input, const std::vector>& outputs, bool is_tensorrt_supported = true, - bool expect_failure = false, const std::string& err_msg = {}) { - OpTester test("Split"); + bool expect_failure = false, bool split_as_input = false, + bool is_initializer = true, const std::string& err_msg = {}) { + int opset_version = split_as_input ? 13 : 7; + OpTester test("Split", opset_version, onnxruntime::kOnnxDomain); test.AddAttribute("axis", axis); - - if (!split_sizes.empty()) - test.AddAttribute("split", split_sizes); - test.AddInput("input", input.first, input.second); + if (!split_sizes.empty()){ + if (split_as_input) + test.AddInput("split", {static_cast(split_sizes.size())}, split_sizes, is_initializer); + else + test.AddAttribute("split", split_sizes); + } int i = 0; for (auto& output : outputs) { @@ -414,7 +418,7 @@ TEST(SplitOperatorTest, InvalidAxis) { outputs.push_back({{1}, {0.f}}); - RunTest(axis, {}, input, outputs, true, true, "Invalid value of attribute 'axis'"); + RunTest(axis, {}, input, outputs, true, true, false, true, "Invalid value of attribute 'axis'"); } // sum of values in splits is too small @@ -434,7 +438,7 @@ TEST(SplitOperatorTest, SplitAttributeSumTooSmall) { outputs.push_back({{1, 2}, {1.f, 2.f}}); outputs.push_back({{2, 2}, {3.f, 4.f, 5.f, 6.f}}); - RunTest(axis, splits, input, outputs, false, true, "Cannot split using values in 'split' attribute"); //TensorRT parser: Assertion failed: axis != BATCH_DIM + RunTest(axis, splits, input, outputs, false, true, false,true, "Cannot split using values in 'split' attribute"); //TensorRT parser: Assertion failed: axis != BATCH_DIM } TEST(SplitOperatorTest, InvalidValueInSplitAttribute) { @@ -452,7 +456,55 @@ TEST(SplitOperatorTest, InvalidValueInSplitAttribute) { outputs.push_back({{1, 2}, {1.f, 2.f}}); outputs.push_back({{3, 2}, {3.f, 4.f, 5.f, 6.f, 7.f, 8.f}}); - RunTest(axis, splits, input, outputs, false, true, "in 'split' attribute"); //TensorRT parser: Assertion failed: axis != BATCH_DIM + RunTest(axis, splits, input, outputs, false, true, false,true, "in 'split' attribute"); //TensorRT parser: Assertion failed: axis != BATCH_DIM +} + +// split as input +TEST(SplitOperatorTest, Axis0UnequalSplitInputFloat) { + const int64_t axis = 0; + std::vector outputs; + + // input shape and data + ShapeAndFloatData input = {{4, 2}, // shape + {1.f, 2.f, + 3.f, 4.f, + 5.f, 6.f, + 7.f, 8.f}}; + + std::vector splits{1, 3}; + + outputs.push_back({{1, 2}, {1.f, 2.f}}); + + outputs.push_back({{3, 2}, + {3.f, 4.f, + 5.f, 6.f, + 7.f, 8.f}}); + + RunTest(axis, splits, input, outputs, false, false, true); +} + +// split as input +TEST(SplitOperatorTest, Axis0UnequalSplitInputFloat_not_initializer) { + const int64_t axis = 0; + std::vector outputs; + + // input shape and data + ShapeAndFloatData input = {{4, 2}, // shape + {1.f, 2.f, + 3.f, 4.f, + 5.f, 6.f, + 7.f, 8.f}}; + + std::vector splits{1, 3}; + + outputs.push_back({{1, 2}, {1.f, 2.f}}); + + outputs.push_back({{3, 2}, + {3.f, 4.f, + 5.f, 6.f, + 7.f, 8.f}}); + + RunTest(axis, splits, input, outputs, false, false, true, false); } /* diff --git a/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc b/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc index 9d7efb0586..20d041032a 100644 --- a/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc @@ -108,5 +108,40 @@ TEST(SqueezeOpTest, SqueezeNegAxis_2) { test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kNGraphExecutionProvider, kOpenVINOExecutionProvider}); } +TEST(SqueezeOpTest, Squeeze_2_axes_input) { + OpTester test("Squeeze", 13); + test.AddInput("data", {1, 4, 1, 1, 2}, + std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); + test.AddInput("axes", {3}, std::vector{0, 2, 3}); + test.AddOutput("squeezed", {4, 2}, + std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); + // Incorrect precision for OpenVINO EP. Will be re-enabled after it's fixed + // TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider, kTensorrtExecutionProvider}); +} + +TEST(SqueezeOpTest, Squeeze_Empty_Axes_opset13) { + OpTester test("Squeeze", 13); + test.AddInput("data", {1, 1, 4, 1}, std::vector(4, 1.0f)); + test.AddOutput("squeezed", {4}, std::vector(4, 1.0f)); + // TensorRT doesn't seem to support missing 'axes' + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider}); +} + +TEST(SqueezeOpTest, SqueezeNegAxis_axes_input) { + OpTester test("Squeeze", 13); + test.AddInput("data", {1, 4, 1, 1, 2}, + std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); + + test.AddInput("axes",{3} ,std::vector{0, -3, -2}); + test.AddOutput("squeezed", {4, 2}, + std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); + + // nGraph does not support neg axis. + // OpenVINO EP Incorrect precision. Will be re-enabled after its fixed + // TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kNGraphExecutionProvider, kOpenVINOExecutionProvider, kTensorrtExecutionProvider}); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/providers/cpu/tensor/unsqueeze_op_test.cc b/onnxruntime/test/providers/cpu/tensor/unsqueeze_op_test.cc index 835451ec8c..f0ce05015f 100644 --- a/onnxruntime/test/providers/cpu/tensor/unsqueeze_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/unsqueeze_op_test.cc @@ -46,42 +46,88 @@ TEST(TensorOpTest, Unsqueeze_3) { } TEST(TensorOpTest, Unsqueeze_Duplicate) { - // This test is valid for opset 12. - // setting opset to -1 makes the test infra pick the latest available opset, so ort will pull the - // schema for that opset and do verification against it. Since opset 13 is enabled this test - //will fail schema validation. - OpTester test("Unsqueeze", 12); + { + OpTester test("Unsqueeze", 12); // opset 1-12 has axes attribute - test.AddAttribute("axes", std::vector{2, 1, 0, 2}); - test.AddInput("input", {2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); - test.AddOutput("output", {1, 1, 1, 2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); - test.Run(OpTester::ExpectResult::kExpectFailure, - "[ShapeInferenceError] 'axes' attribute must not contain any duplicates", - {kTensorrtExecutionProvider}); //TensorRT failed + test.AddAttribute("axes", std::vector{2, 1, 0, 2}); + test.AddInput("input", {2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + test.AddOutput("output", {1, 1, 1, 2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + test.Run(OpTester::ExpectResult::kExpectFailure, + "[ShapeInferenceError] 'axes' attribute must not contain any duplicates", + {kTensorrtExecutionProvider}); //TensorRT failed + } + { + OpTester test("Unsqueeze", -1); // use latest opset with axis input + + test.AddInput("input", {2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + test.AddInput("axes", {4}, std::vector{2, 1, 0, 2}, true); //set as initializer to enable shape inference + test.AddOutput("output", {1, 1, 1, 2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + test.Run(OpTester::ExpectResult::kExpectFailure, + "[ShapeInferenceError] 'axes' attribute must not contain any duplicates", + {kTensorrtExecutionProvider}); //TensorRT failed + } } TEST(TensorOpTest, Unsqueeze_OutOfRange) { - // This test is valid for opset 12. - // setting opset to -1 makes the test infra pick the latest available opset, so ort will pull the - // schema for that opset and do verification against it. Since opset 13 is enabled this test - // will fail schema validation. - OpTester test("Unsqueeze", 12); - - test.AddAttribute("axes", std::vector{4}); - test.AddInput("input", {2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); - test.AddOutput("output", {2, 1, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); - test.Run(OpTester::ExpectResult::kExpectFailure, - "[ShapeInferenceError] values in 'axes' are beyond the bounds of the computed output shape"); + { + OpTester test("Unsqueeze", 12); // opset 1-12 has axes attribute + test.AddAttribute("axes", std::vector{4}); + test.AddInput("input", {2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + test.AddOutput("output", {2, 1, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + test.Run(OpTester::ExpectResult::kExpectFailure, + "[ShapeInferenceError] values in 'axes' are beyond the bounds of the computed output shape"); + } + { + OpTester test("Unsqueeze", -1); // use latest opset with axis input + test.AddInput("input", {2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + test.AddInput("axes", {1}, std::vector{4}, true); //set as initializer to enable shape inference + test.AddOutput("output", {2, 1, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + // nGraph and TensorRT does not support negative axis. + test.Run(OpTester::ExpectResult::kExpectFailure, + "[ShapeInferenceError] values in 'axes' are beyond the bounds of the computed output shape", + {kTensorrtExecutionProvider}); //TensorRT expects 'axes' attribute + } } TEST(TensorOpTest, UnsqueezeNegAxis_3) { - OpTester test("Unsqueeze", 12); // use latest opset for shape inference errors + { + OpTester test("Unsqueeze", 12); // opset 1-12 has axes attribute + test.AddAttribute("axes", std::vector{-4, 1, -6}); + test.AddInput("input", {2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + test.AddOutput("output", {1, 1, 1, 2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + // nGraph and TensorRT does not support negative axis. + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kNGraphExecutionProvider, kTensorrtExecutionProvider}); + } + { + OpTester test("Unsqueeze", 13); // use latest opset with axis input + test.AddInput("input", {2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + test.AddInput("axes", {3}, std::vector{-4, 1, -6}); + test.AddOutput("output", {1, 1, 1, 2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + // nGraph and TensorRT does not support negative axis. + // TODO: TensorRT, OpenVINO dont support "axes" input in opset 13, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kNGraphExecutionProvider, + kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); + } +} + +TEST(TensorOpTest, Unsqueeze_1_int32_axes_input) { + OpTester test("Unsqueeze", 13); + + test.AddInput("input", {2, 3, 4}, std::vector(2 * 3 * 4, 1)); + test.AddInput("axes", {1}, std::vector{1}); + test.AddOutput("output", {2, 1, 3, 4}, std::vector(2 * 3 * 4, 1)); + // TODO: TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); +} + +TEST(TensorOpTest, Unsqueeze_3_axes_input) { + OpTester test("Unsqueeze", 13); - test.AddAttribute("axes", std::vector{-4, 1, -6}); test.AddInput("input", {2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); + test.AddInput("axes", {3}, std::vector{2, 1, 0}); test.AddOutput("output", {1, 1, 1, 2, 3, 4}, std::vector(2 * 3 * 4, 1.0f)); - // nGraph and TensorRT does not support negative axis. - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kNGraphExecutionProvider, kTensorrtExecutionProvider}); + // TODO: TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); } } // namespace test diff --git a/orttraining/orttraining/core/graph/gradient_builder.cc b/orttraining/orttraining/core/graph/gradient_builder.cc index 9ba5655493..778c3a28c0 100644 --- a/orttraining/orttraining/core/graph/gradient_builder.cc +++ b/orttraining/orttraining/core/graph/gradient_builder.cc @@ -792,11 +792,19 @@ IMPLEMENT_GRADIENT_BUILDER(GetLogSoftmaxGradient) { } IMPLEMENT_GRADIENT_BUILDER(GetUnsqueezeGradient) { - return std::vector{ - NodeDef("Squeeze", - {GO(0)}, - {GI(0)}, - SrcNodeAttributes())}; + if (SrcNodeOpsetVersion() < 13) { + return std::vector{ + NodeDef("Squeeze", + {GO(0)}, + {GI(0)}, + SrcNodeAttributes())}; + } else { // mandatory input 'axes' since opset 13 + return std::vector{ + NodeDef(OpDef{"Squeeze", kOnnxDomain, 13}, + {GO(0), I(1)}, + {GI(0)}, + SrcNodeAttributes())}; + } } IMPLEMENT_GRADIENT_BUILDER(GetGatherGradient) { @@ -830,26 +838,34 @@ IMPLEMENT_GRADIENT_BUILDER(GetReluGradient) { IMPLEMENT_GRADIENT_BUILDER(GetSqueezeGradient) { std::vector result; - auto attributes = SrcNodeAttributes(); - std::vector axes_values; - if (attributes.find("axes") != attributes.end()) { - axes_values = RetrieveValues(attributes.at("axes")); + size_t numInputs = GetSrcNodeInputSize(); + if (SrcNodeOpsetVersion() < 13) { //axes attribute + auto attributes = SrcNodeAttributes(); + std::vector axes_values; + if (attributes.find("axes") != attributes.end()) { + axes_values = RetrieveValues(attributes.at("axes")); + result.push_back( + NodeDef("Unsqueeze", + {GO(0)}, + {GI(0)}, + {MakeAttribute("axes", axes_values)})); + } + } else if(numInputs == 2){ //optional input 'axes' is provided result.push_back( - NodeDef("Unsqueeze", - {GO(0)}, - {GI(0)}, - {MakeAttribute("axes", axes_values)})); - // if axes attribute not provided for squeeze - } else { - result.push_back( - NodeDef("Shape", - {I(0)}, - {IA("I0_shape")})); - result.push_back( - NodeDef("Reshape", - {GO(0), IA("I0_shape")}, + 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)})); } + return result; } @@ -1119,10 +1135,20 @@ IMPLEMENT_GRADIENT_BUILDER(GetReduceSumGradient) { } ArgDef grad = GO(0); - if (!keepdims && attributes.find("axes") != attributes.end()) { - std::vector axes_values = RetrieveValues(attributes.at("axes")); - grad = IA("Unqueezed_Grad"); - result.push_back(NodeDef("Unsqueeze", {GO(0)}, {grad}, {MakeAttribute("axes", axes_values)})); + if (!keepdims) { + size_t numInputs = GetSrcNodeInputSize(); + if (SrcNodeOpsetVersion() < 13) { //axes is attribute + if (attributes.find("axes") != attributes.end()) { + std::vector axes_values = RetrieveValues(attributes.at("axes")); + + 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 } result.push_back(NodeDef("Shape", {I(0)}, {IA("Shaped_X")})); diff --git a/orttraining/orttraining/core/graph/gradient_builder_base.h b/orttraining/orttraining/core/graph/gradient_builder_base.h index 7822d1ec11..d88f28eda7 100644 --- a/orttraining/orttraining/core/graph/gradient_builder_base.h +++ b/orttraining/orttraining/core/graph/gradient_builder_base.h @@ -181,6 +181,10 @@ class GradientBuilderBase { return node_->OpType(); } + int SrcNodeOpsetVersion() const { + return node_->Op()->since_version(); + } + template static NodeDef ConstantVectorNode(const std::vector& values, const std::string& arg_name) { auto t_proto = ONNX_NAMESPACE::ToTensor(values); diff --git a/orttraining/orttraining/test/gradient/gradient_ops_test.cc b/orttraining/orttraining/test/gradient/gradient_ops_test.cc index b8fa8a2a0e..07799b418e 100644 --- a/orttraining/orttraining/test/gradient/gradient_ops_test.cc +++ b/orttraining/orttraining/test/gradient/gradient_ops_test.cc @@ -38,63 +38,78 @@ static bool IsErrorWithinTolerance(float error, float tolerance) { #define EXPECT_IS_TINY(max_error) \ EXPECT_IS_TINIER_THAN(max_error, 1.5e-2f) -static void RunReductionTests(const OpDef& op_def) { - TestDataVector test_data( - // Input X - { - {{4, 3, 2}}, - {{4, 3, 2}}, - {{4, 3, 2}}, - {{4, 3, 2}}, - {{4, 3, 2}}, - {{4, 3, 2}}, - {{4, 3, 2}}, - {{4, 3, 2}}, - }, - // Input Y - { - {{1, 1, 1}}, - {{}}, - {{1, 3, 1}}, - {{2}}, - {{4, 1, 2}}, - {{4, 3}}, - {{4, 1, 2}}, - {{4}}}, - // Attributes - { - // default - {}, - // axes = [0, 1, 2], keepdims = 0 - {MakeAttribute("axes", std::vector{0, 1, 2}), - MakeAttribute("keepdims", int64_t(0))}, - // axes = [0, 2], keepdims = 1 - {MakeAttribute("axes", std::vector{0, 2})}, - // axes = [0, 1], keepdims = 0 - {MakeAttribute("axes", std::vector{0, 1}), - MakeAttribute("keepdims", int64_t(0))}, - // axes = [1], keepdims = 1 - {MakeAttribute("axes", std::vector{1}), - MakeAttribute("keepdims", int64_t(1))}, - // axes = [2], keepdims = 0 - {MakeAttribute("axes", std::vector{2}), - MakeAttribute("keepdims", int64_t(0))}, - // axes = [-2], keepdims = 1 - {MakeAttribute("axes", std::vector{-2}), - MakeAttribute("keepdims", int64_t(1))}, - // axes = [-2, -1], keepdims = 0 - {MakeAttribute("axes", std::vector{-2, -1}), - MakeAttribute("keepdims", int64_t(0))}}); +static void RunReductionTests(const OpDef& op_def, + bool axes_as_input = false, + bool check_not_have_shape_inferencing = false) { + std::vector> + x_shapes = { + {4, 3, 2}, + {4, 3, 2}, + {4, 3, 2}, + {4, 3, 2}, + {4, 3, 2}, + {4, 3, 2}, + {4, 3, 2}, + {4, 3, 2}, + }; + std::vector> y_shapes = { + {1, 1, 1}, + {}, + {1, 3, 1}, + {2}, + {4, 1, 2}, + {4, 3}, + {4, 1, 2}, + {4}, + }; + std::vector> axes_vec = { + {}, //default case + {0, 1, 2}, + {0, 2}, + {0, 1}, + {1}, + {2}, + {-2}, + {-2, -1}, + }; + std::vector keepdims_ip = { + -1, //default case + 0, + 1, + 0, + 1, + 0, + 1, + 0, + }; GradientChecker gradient_checker; float max_error; - - for (size_t i = 0; i < std::get<0>(test_data).size(); i++) { + for (size_t i = 0; i < x_shapes.size(); i++) { max_error = 0; - gradient_checker.ComputeGradientError(op_def, std::get<0>(test_data)[i], - std::get<1>(test_data)[i], &max_error, - std::get<2>(test_data)[i]); + TensorShape x_shape = x_shapes[i]; + TensorShape y_shape = y_shapes[i]; + std::vector axes = axes_vec[i]; + std::vector> x_datas; + RandomValueGenerator random{}; + x_datas.push_back(random.Gaussian(x_shapes[i], 0.f, 5.f)); + std::vector input = {x_shape}; + std::vector attributes = {}; + if (keepdims_ip[i] != -1) attributes.push_back(MakeAttribute("keepdims", keepdims_ip[i])); + if (axes_as_input) { + std::vector axes_float; + std::transform(begin(axes), end(axes), std::back_inserter(axes_float), [](int64_t i) { return static_cast(i); }); + TensorInfo axes_info({static_cast(axes.size())}, false, nullptr, DataTypeImpl::GetTensorType()); + input.push_back(axes_info); + x_datas.push_back(axes_float); + } else { + if (axes.size() > 0) + attributes.push_back(MakeAttribute("axes", axes)); + } + + gradient_checker.ComputeGradientError(op_def, input, {y_shape}, &max_error, x_datas, + attributes, true, check_not_have_shape_inferencing); EXPECT_IS_TINY(max_error); } } @@ -399,7 +414,7 @@ TEST(GradientCheckerTest, NegGrad) { } TEST(GradientCheckerTest, LogGrad) { - TensorShape shape({2,5,6}); + TensorShape shape({2, 5, 6}); std::function transformer = [](float x) { return std::fabs(x) + 1e-1f; }; TensorInfo x_info{shape, true, &transformer}; @@ -548,9 +563,14 @@ TEST(GradientCheckerTest, ReduceMeanGrad) { TEST(GradientCheckerTest, ReduceSumGrad) { // Attribute axes supports negative values from opset 11. - OpDef op_def{"ReduceSum", kOnnxDomain, 11}; + OpDef op_def_11{"ReduceSum", kOnnxDomain, 11}; - RunReductionTests(op_def); + RunReductionTests(op_def_11, false, true); + + // axes is input from opset 13. + OpDef op_def_13{"ReduceSum", kOnnxDomain, 13}; + + RunReductionTests(op_def_13, true, true); } TEST(GradientCheckerTest, ReduceLogSumExpGrad) { @@ -590,6 +610,12 @@ TEST(GradientCheckerTest, SplitGrad) { gradient_checker.ComputeGradientError(op_def, {shape}, {{3, 5}, {3, 5}, {3, 5}}, &max_error, {MakeAttribute("axis", int64_t(0))}); EXPECT_IS_TINY(max_error); + + //opset13 test + OpDef op_def_13{"Split", kOnnxDomain, 13}; + gradient_checker.ComputeGradientError(op_def_13, {shape}, {{3, 5}, {3, 5}, {3, 5}}, &max_error, + {MakeAttribute("axis", int64_t(0))}); + EXPECT_IS_TINY(max_error); } template @@ -752,9 +778,9 @@ TEST(GradientCheckerTest, ConvGrad) { } static void TestConcatOpGrad(const std::string& op_type, - const std::string& domain = kOnnxDomain, - int opset_version = 9, - bool check_not_have_shape_inferencing = false) { + const std::string& domain = kOnnxDomain, + int opset_version = 9, + bool check_not_have_shape_inferencing = false) { float max_error; GradientChecker gradient_checker; const bool extra_input = op_type == "ConcatTraining"; @@ -767,9 +793,9 @@ static void TestConcatOpGrad(const std::string& op_type, std::vector output = {y_shape}; if (extra_input) output.push_back(TensorInfo({3}, false, nullptr, DataTypeImpl::GetTensorType())); gradient_checker.ComputeGradientError(op_def, {x_shape, x_shape, x_shape}, - output, &max_error, - {MakeAttribute("axis", int64_t(0))}, true, - check_not_have_shape_inferencing); + output, &max_error, + {MakeAttribute("axis", int64_t(0))}, true, + check_not_have_shape_inferencing); EXPECT_IS_TINY(max_error); } @@ -781,7 +807,7 @@ static void TestConcatOpGrad(const std::string& op_type, if (extra_input) output.push_back(TensorInfo({3}, false, nullptr, DataTypeImpl::GetTensorType())); gradient_checker.ComputeGradientError(op_def, {x_shape, x_shape, x_shape}, output, &max_error, - {MakeAttribute("axis", int64_t(1))}, true, + {MakeAttribute("axis", int64_t(1))}, true, check_not_have_shape_inferencing); EXPECT_IS_TINY(max_error); } @@ -794,7 +820,7 @@ static void TestConcatOpGrad(const std::string& op_type, if (extra_input) output.push_back(TensorInfo({3}, false, nullptr, DataTypeImpl::GetTensorType())); gradient_checker.ComputeGradientError(op_def, {x_shape, x_shape, x_shape}, output, &max_error, - {MakeAttribute("axis", int64_t(2))}, true, + {MakeAttribute("axis", int64_t(2))}, true, check_not_have_shape_inferencing); EXPECT_IS_TINY(max_error); } @@ -808,7 +834,7 @@ static void TestConcatOpGrad(const std::string& op_type, if (extra_input) output.push_back(TensorInfo({2}, false, nullptr, DataTypeImpl::GetTensorType())); gradient_checker.ComputeGradientError(op_def, {x1_shape, x2_shape}, output, &max_error, - {MakeAttribute("axis", int64_t(1))}, true, + {MakeAttribute("axis", int64_t(1))}, true, check_not_have_shape_inferencing); EXPECT_IS_TINY(max_error); } @@ -822,7 +848,7 @@ static void TestConcatOpGrad(const std::string& op_type, if (extra_input) output.push_back(TensorInfo({2}, false, nullptr, DataTypeImpl::GetTensorType())); gradient_checker.ComputeGradientError(op_def, {x1_shape, x2_shape}, output, &max_error, - {MakeAttribute("axis", int64_t(-1))}, true, + {MakeAttribute("axis", int64_t(-1))}, true, check_not_have_shape_inferencing); EXPECT_IS_TINY(max_error); } @@ -985,90 +1011,97 @@ TEST(GradientCheckerTest, TransposeGrad) { } } -TEST(GradientCheckerTest, UnsqueezeGrad) { +static void RunSqueezeUnsqueezeTests(const OpDef& op_def, + std::vector> x_shapes, + std::vector> y_shapes, + std::vector> axes_ip, + bool axes_input = false) { float max_error; GradientChecker gradient_checker; - OpDef op_def{"Unsqueeze"}; float error_tolerance = 1e-3f; - { - TensorShape x_shape({2, 3}); - TensorShape y_shape({1, 2, 3, 1}); - std::vector axes{0, 3}; - gradient_checker.ComputeGradientError(op_def, {x_shape}, {y_shape}, &max_error, - {MakeAttribute("axes", axes)}); - EXPECT_IS_TINIER_THAN(max_error, error_tolerance); - } + for (size_t i = 0; i < x_shapes.size(); i++) { + TensorShape x_shape = x_shapes[i]; + TensorShape y_shape = y_shapes[i]; + std::vector axes = axes_ip[i]; + std::vector> x_datas; + RandomValueGenerator random{}; + x_datas.push_back(random.Gaussian(x_shapes[i], 0.f, 5.f)); + std::vector input = {x_shape}; + std::vector attributes = {}; + if (axes_input) { + std::vector axes_float; + std::transform(begin(axes), end(axes), std::back_inserter(axes_float), [](int64_t i) { return static_cast(i); }); + TensorInfo axes_info({static_cast(axes.size())}, false, nullptr, DataTypeImpl::GetTensorType()); + input.push_back(axes_info); + x_datas.push_back(axes_float); + } else { + attributes.push_back(MakeAttribute("axes", axes)); + } - { - TensorShape x_shape({2, 3}); - TensorShape y_shape({1, 1, 2, 3}); - std::vector axes{0, 1}; - gradient_checker.ComputeGradientError(op_def, {x_shape}, {y_shape}, &max_error, - {MakeAttribute("axes", axes)}); - EXPECT_IS_TINIER_THAN(max_error, error_tolerance); - } - - { - TensorShape x_shape({2, 3}); - TensorShape y_shape({1, 2, 1, 3, 1}); - std::vector axes{0, 2, 4}; - gradient_checker.ComputeGradientError(op_def, {x_shape}, {y_shape}, &max_error, - {MakeAttribute("axes", axes)}); + gradient_checker.ComputeGradientError(op_def, input, {y_shape}, &max_error, x_datas, attributes); EXPECT_IS_TINIER_THAN(max_error, error_tolerance); } } TEST(GradientCheckerTest, SqueezeGrad) { - float max_error; - GradientChecker gradient_checker; - OpDef op_def{"Squeeze"}; - float error_tolerance = 1e-3f; - - { - TensorShape x_shape({1, 2, 3, 1}); - TensorShape y_shape({2, 3}); - std::vector axes{0, 3}; - gradient_checker.ComputeGradientError(op_def, {x_shape}, {y_shape}, &max_error, - {MakeAttribute("axes", axes)}); - EXPECT_IS_TINIER_THAN(max_error, error_tolerance); - } - - { - TensorShape x_shape({1, 1, 2, 3, 4}); - TensorShape y_shape({2, 3, 4}); - std::vector axes{0, 1}; - gradient_checker.ComputeGradientError(op_def, {x_shape}, {y_shape}, &max_error, - {MakeAttribute("axes", axes)}); - EXPECT_IS_TINIER_THAN(max_error, error_tolerance); - } - - { - TensorShape x_shape({1, 2, 1, 3, 1}); - TensorShape y_shape({2, 3}); - std::vector axes{0, 2, 4}; - gradient_checker.ComputeGradientError(op_def, {x_shape}, {y_shape}, &max_error, - {MakeAttribute("axes", axes)}); - EXPECT_IS_TINIER_THAN(max_error, error_tolerance); - } - - { - TensorShape x_shape({1, 2, 1, 3, 1}); - TensorShape y_shape({1, 2, 3, 1}); - std::vector axes{2}; - gradient_checker.ComputeGradientError(op_def, {x_shape}, {y_shape}, &max_error, - {MakeAttribute("axes", axes)}); - EXPECT_IS_TINIER_THAN(max_error, error_tolerance); - } - /* TODO: enable test with no axis when squeeze kernel is fixed (separate bug filed) - { TensorShape x_shape({1, 2, 1, 3, 1}); TensorShape y_shape({2, 3}); - gradient_checker.ComputeGradientError(op_def, {x_shape}, {y_shape}, &max_error); - EXPECT_IS_TINIER_THAN(max_error, error_tolerance); - } */ + std::vector> x_shapes = { + {1, 2, 3, 1}, + {1, 1, 2, 3, 4}, + {1, 2, 1, 3, 1}, + {1, 2, 1, 3, 1}, + // {1, 2, 1, 3, 1}, + }; + std::vector> y_shapes = { + {2, 3}, + {2, 3, 4}, + {2, 3}, + {1, 2, 3, 1}, + // {2, 3}, + }; + std::vector> axes_ip = { + {0, 3}, + {0, 1}, + {0, 2, 4}, + {2}, + // {} + }; + + OpDef op_def{"Squeeze"}; + RunSqueezeUnsqueezeTests(op_def, x_shapes, y_shapes, axes_ip); + + //axes as input from opset 13 + OpDef op_def_2{"Squeeze", kOnnxDomain, 13}; + RunSqueezeUnsqueezeTests(op_def_2, x_shapes, y_shapes, axes_ip, true); +} + +TEST(GradientCheckerTest, UnsqueezeGrad) { + std::vector> x_shapes = { + {2, 3}, + {2, 3}, + {2, 3}, + }; + std::vector> y_shapes = { + {1, 2, 3, 1}, + {1, 1, 2, 3}, + {1, 2, 1, 3, 1}, + }; + std::vector> axes_ip = { + {0, 3}, + {0, 1}, + {0, 2, 4}, + }; + + OpDef op_def{"Unsqueeze"}; + RunSqueezeUnsqueezeTests(op_def, x_shapes, y_shapes, axes_ip); + + //axes as input from opset 13 + OpDef op_def_2{"Unsqueeze", kOnnxDomain, 13}; + RunSqueezeUnsqueezeTests(op_def_2, x_shapes, y_shapes, axes_ip, true); } // TODO: Reshape missing @@ -1234,7 +1267,7 @@ void GradientCheckerSoftmaxGradHelper(bool is_log_softmax) { float max_error; GradientChecker gradient_checker; - const std::string op = is_log_softmax? "LogSoftmax" : "Softmax"; + const std::string op = is_log_softmax ? "LogSoftmax" : "Softmax"; OpDef op_def{op}; // default_axis