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

Ryanunderhill/concat neg axis (#152)

Browse source 
* Add negative axis support to concat

* Unit test bug
This commit is contained in:
Ryan Hill 2018-12-12 18:01:28 -08:00 committed by GitHub
parent cfd3b4c606
commit 30ab01e7be
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 51 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
onnxruntime/core/providers/cpu/tensor/concat.cc
View file

@ -2,6 +2,7 @@
// Licensed under the MIT License.
#include "core/providers/cpu/tensor/concat.h"
#include "core/providers/common.h"
namespace onnxruntime {
@ -17,6 +18,8 @@ Status ConcatBase::PrepareForCompute(OpKernelContext* ctx, int input_count, Prep
if (tensor_pointer == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch");
const Tensor& inputs_0 = *tensor_pointer;
auto axis = HandleNegativeAxis(axis_, inputs_0.Shape().NumDimensions());
// Ensure all of the non concatenated axes match each other
for (int index = 1; index < input_count; index++) {
tensor_pointer = ctx->Input<Tensor>(index);
@ -25,7 +28,7 @@ Status ConcatBase::PrepareForCompute(OpKernelContext* ctx, int input_count, Prep
// Ensure all the other axes match
auto dimension_count = inputs_0.Shape().NumDimensions();
for (int axis_index = 0; axis_index < dimension_count; axis_index++) {
if (axis_index == axis_)
if (axis_index == axis)
continue;
ONNXRUNTIME_RETURN_IF_NOT(data_n.Shape()[axis_index] == inputs_0.Shape()[axis_index], "Non concat axis dimensions must match: Axis ", axis_index, " has mismatched dimensions of ", data_n.Shape()[axis_index], " and ", inputs_0.Shape()[axis_index]);
}
@ -36,19 +39,19 @@ Status ConcatBase::PrepareForCompute(OpKernelContext* ctx, int input_count, Prep
for (int index = 0; index < input_count; index++) {
tensor_pointer = ctx->Input<Tensor>(index);
if (tensor_pointer == nullptr) return Status(common::ONNXRUNTIME, common::FAIL, "input count mismatch");
concat_axis_size += tensor_pointer->Shape()[int(axis_)];
concat_axis_size += tensor_pointer->Shape()[int(axis)];
}
// Calculate the shape of the output tensor
std::vector<int64_t> dims;
for (int dimension_index = 0; dimension_index < inputs_0.Shape().NumDimensions(); dimension_index++)
dims.emplace_back(inputs_0.Shape()[dimension_index]);
dims[axis_] = concat_axis_size;
dims[axis] = concat_axis_size;
TensorShape outputShape(dims);
// The output_axis_pitch is the number of elements to add to move to the next split axis in the output
p.output_axis_pitch = 1;
for (auto i = int64_t(dims.size()); i-- > axis_;)
for (auto i = int64_t(dims.size()); i-- > axis;)
p.output_axis_pitch *= dims[i];
auto& concat_result = *ctx->Output(0, outputShape);
@ -63,7 +66,7 @@ Status ConcatBase::PrepareForCompute(OpKernelContext* ctx, int input_count, Prep
// The input_axis_pitch is the number of elements to add to move to the next split axis in the input
int64_t input_axis_pitch = 1;
for (int i = int(data_n.Shape().NumDimensions()); i-- > axis_;)
for (int i = int(data_n.Shape().NumDimensions()); i-- > axis;)
input_axis_pitch *= data_n.Shape()[i];
p.inputs.push_back({&data_n, input_axis_pitch});

43
onnxruntime/test/providers/cpu/tensor/concat_op_test.cc
View file

@ -29,6 +29,17 @@ TEST(MathOpTest, Concat1D_int32) {
test.Run();
}
TEST(MathOpTest, Concat1D_int32_negative_axis) {
OpTester test("Concat");
test.AddAttribute("axis", int64_t{-1});
test.AddInput<int32_t>("input1", {1}, {1});
test.AddInput<int32_t>("input2", {2}, {2, 3});
test.AddInput<int32_t>("input3", {4}, {4, 5, 6, 7});
test.AddOutput<int32_t>("concat_result", {7}, {1, 2, 3, 4, 5, 6, 7});
test.Run();
}
TEST(MathOpTest, Concat1D) {
OpTester test("Concat");
test.AddAttribute("axis", int64_t{0});
@ -103,6 +114,38 @@ TEST(MathOpTest, Concat3D_1) {
test.Run();
}
TEST(MathOpTest, Concat3D_1_negative_axis) {
OpTester test("Concat");
test.AddAttribute("axis", int64_t{-3});
std::vector<int64_t> dims{1, 3, 3};
test.AddInput<float>("input1", dims,
{111.0f, 112.0f, 113.0f,
121.0f, 122.0f, 123.0f,
131.0f, 132.0f, 133.0f});
test.AddInput<float>("input2", dims,
{211.0f, 212.0f, 213.0f,
221.0f, 222.0f, 223.0f,
231.0f, 232.0f, 233.0f});
test.AddInput<float>("input3", dims,
{311.0f, 312.0f, 313.0f,
321.0f, 322.0f, 323.0f,
331.0f, 332.0f, 333.0f});
test.AddOutput<float>("concat_result", {3, 3, 3},
{111.0f, 112.0f, 113.0f,
121.0f, 122.0f, 123.0f,
131.0f, 132.0f, 133.0f,
211.0f, 212.0f, 213.0f,
221.0f, 222.0f, 223.0f,
231.0f, 232.0f, 233.0f,
311.0f, 312.0f, 313.0f,
321.0f, 322.0f, 323.0f,
331.0f, 332.0f, 333.0f});
test.Run();
}
TEST(MathOpTest, Concat3D_2) {
OpTester test("Concat");
test.AddAttribute("axis", int64_t{1});