diff --git a/onnxruntime/core/providers/cpu/tensor/upsample.cc b/onnxruntime/core/providers/cpu/tensor/upsample.cc index 4b2c5f6c57..4e9871af43 100644 --- a/onnxruntime/core/providers/cpu/tensor/upsample.cc +++ b/onnxruntime/core/providers/cpu/tensor/upsample.cc @@ -96,52 +96,66 @@ Status UpsampleNearest(const T* input, int64_t output_idx = 0; int64_t input_idx = 0; -#define OneDimensionProcessor(dim_inx) \ - use_extrapolation_value[dim_inx] = false; \ - float original_##dim_inx##_idx = get_original_coordinate(static_cast(output_dim##dim_inx##_inx), scales[dim_inx], static_cast(output_shape[dim_inx]), static_cast(input_shape[dim_inx]), roi[dim_inx], roi[n_dim + dim_inx]); \ - if (extrapolation_enabled && (original_##dim_inx##_idx < 0 || original_##dim_inx##_idx > input_shape[dim_inx] - 1)) use_extrapolation_value[dim_inx] = true; \ - int64_t input_dim##dim_inx##_inx = get_nearest_pixel(original_##dim_inx##_idx, scales[dim_inx] < 1); \ - if (input_dim##dim_inx##_inx > input_shape[dim_inx] - 1) input_dim##dim_inx##_inx = input_shape[dim_inx] - 1; \ - if (input_dim##dim_inx##_inx < 0) input_dim##dim_inx##_inx = 0; \ - if (input_dim##dim_inx##_inx != input_dim_counters[dim_inx]) { \ - input_idx += (input_dim##dim_inx##_inx - input_dim_counters[dim_inx]) * input_dim_factor[dim_inx]; \ - input_dim_counters[dim_inx] = input_dim##dim_inx##_inx; \ - } - if (n_dim == 1) { for (int64_t output_dim0_inx = 0; output_dim0_inx < output_shape[0]; output_dim0_inx++) { - OneDimensionProcessor(0); - output[output_idx++] = use_extrapolation_value[0] ? static_cast(extrapolation_value) : input[input_idx]; + use_extrapolation_value[0] = false; + float original_0_idx = get_original_coordinate(static_cast(output_dim0_inx), scales[0], + static_cast(output_shape[0]), static_cast(input_shape[0]), + roi[0], roi[n_dim + 0]); + if (extrapolation_enabled && (original_0_idx < 0 || original_0_idx > input_shape[0] - 1)) use_extrapolation_value[0] = true; + int64_t input_dim0_inx = get_nearest_pixel(original_0_idx, scales[0] < 1); + if (input_dim0_inx > input_shape[0] - 1) input_dim0_inx = input_shape[0] - 1; + if (input_dim0_inx < 0) input_dim0_inx = 0; + output[output_idx++] = use_extrapolation_value[0] ? static_cast(extrapolation_value) : input[input_dim0_inx]; } return Status::OK(); } + auto CalculateInputMapping = + [n_dim, &input_shape, &output_shape, &input_dim_factor, &scales, &roi, extrapolation_enabled, &get_original_coordinate, &get_nearest_pixel]( + std::vector& input_mapping, const int64_t axis) { + const int64_t input_size = input_dim_factor[0] * input_shape[0]; + for (int64_t dim = 0; dim < output_shape[axis]; dim++) { + float original_dim = get_original_coordinate(static_cast(dim), scales[axis], static_cast(output_shape[axis]), + static_cast(input_shape[axis]), roi[axis], roi[n_dim + axis]); + bool need_extrapolation = (extrapolation_enabled && (original_dim < 0 || original_dim > input_shape[axis] - 1)); + int64_t input_dim = get_nearest_pixel(original_dim, scales[axis] < 1); + if (input_dim >= input_shape[axis]) input_dim = input_shape[axis] - 1; + if (input_dim < 0) input_dim = 0; + input_mapping[dim] = need_extrapolation ? (-input_size) : (input_dim * input_dim_factor[axis]); + } + return; + }; + if (n_dim == 2) { + std::vector input_mapping_0(output_shape[0]); + std::vector input_mapping_1(output_shape[1]); + CalculateInputMapping(input_mapping_0, 0); + CalculateInputMapping(input_mapping_1, 1); for (int64_t output_dim0_inx = 0; output_dim0_inx < output_shape[0]; output_dim0_inx++) { - OneDimensionProcessor(0); + int64_t input_idx_0 = input_mapping_0[output_dim0_inx]; for (int64_t output_dim1_inx = 0; output_dim1_inx < output_shape[1]; output_dim1_inx++) { - OneDimensionProcessor(1); - output[output_idx++] = (use_extrapolation_value[0] || use_extrapolation_value[1]) - ? static_cast(extrapolation_value) - : input[input_idx]; + int64_t input_idx_1 = input_idx_0 + input_mapping_1[output_dim1_inx]; + output[output_idx++] = (input_idx_1 < 0) ? static_cast(extrapolation_value) : input[input_idx_1]; } } return Status::OK(); } if (n_dim == 3) { + std::vector input_mapping_0(output_shape[0]); + std::vector input_mapping_1(output_shape[1]); + std::vector input_mapping_2(output_shape[2]); + CalculateInputMapping(input_mapping_0, 0); + CalculateInputMapping(input_mapping_1, 1); + CalculateInputMapping(input_mapping_2, 2); for (int64_t output_dim0_inx = 0; output_dim0_inx < output_shape[0]; output_dim0_inx++) { - OneDimensionProcessor(0); + int64_t input_idx_0 = input_mapping_0[output_dim0_inx]; for (int64_t output_dim1_inx = 0; output_dim1_inx < output_shape[1]; output_dim1_inx++) { - OneDimensionProcessor(1); + int64_t input_idx_1 = input_idx_0 + input_mapping_1[output_dim1_inx]; for (int64_t output_dim2_inx = 0; output_dim2_inx < output_shape[2]; output_dim2_inx++) { - OneDimensionProcessor(2); - bool use_extrapolation = std::any_of(use_extrapolation_value.begin(), use_extrapolation_value.end(), - [](bool use_extrapolation) { - return use_extrapolation == true; - }); - - output[output_idx++] = use_extrapolation ? static_cast(extrapolation_value) : input[input_idx]; + int64_t input_idx_2 = input_idx_1 + input_mapping_2[output_dim2_inx]; + output[output_idx++] = (input_idx_2 < 0) ? static_cast(extrapolation_value) : input[input_idx_2]; } } } @@ -153,20 +167,23 @@ Status UpsampleNearest(const T* input, UpsampleNearest2x(input_shape[0], input_shape[1], input_shape[2], input_shape[3], input, output); return Status::OK(); } + std::vector input_mapping_0(output_shape[0]); + std::vector input_mapping_1(output_shape[1]); + std::vector input_mapping_2(output_shape[2]); + std::vector input_mapping_3(output_shape[3]); + CalculateInputMapping(input_mapping_0, 0); + CalculateInputMapping(input_mapping_1, 1); + CalculateInputMapping(input_mapping_2, 2); + CalculateInputMapping(input_mapping_3, 3); for (int64_t output_dim0_inx = 0; output_dim0_inx < output_shape[0]; output_dim0_inx++) { - OneDimensionProcessor(0); + int64_t input_idx_0 = input_mapping_0[output_dim0_inx]; for (int64_t output_dim1_inx = 0; output_dim1_inx < output_shape[1]; output_dim1_inx++) { - OneDimensionProcessor(1); + int64_t input_idx_1 = input_idx_0 + input_mapping_1[output_dim1_inx]; for (int64_t output_dim2_inx = 0; output_dim2_inx < output_shape[2]; output_dim2_inx++) { - OneDimensionProcessor(2); + int64_t input_idx_2 = input_idx_1 + input_mapping_2[output_dim2_inx]; for (int64_t output_dim3_inx = 0; output_dim3_inx < output_shape[3]; output_dim3_inx++) { - OneDimensionProcessor(3); - bool use_extrapolation = std::any_of(use_extrapolation_value.begin(), use_extrapolation_value.end(), - [](bool use_extrapolation) { - return use_extrapolation == true; - }); - - output[output_idx++] = use_extrapolation ? static_cast(extrapolation_value) : input[input_idx]; + int64_t input_idx_3 = input_idx_2 + input_mapping_3[output_dim3_inx]; + output[output_idx++] = (input_idx_3 < 0) ? static_cast(extrapolation_value) : input[input_idx_3]; } } } @@ -174,35 +191,28 @@ Status UpsampleNearest(const T* input, return Status::OK(); } -#undef OneDimensionProcessor + std::vector> input_mappings(n_dim); + for (int64_t dim_idx = 0; dim_idx < n_dim; ++dim_idx) { + input_mappings[dim_idx].resize(output_shape[dim_idx]); + CalculateInputMapping(input_mappings[dim_idx], dim_idx); + } std::vector output_dim_counter(n_dim); - output_dim_counter[n_dim - 1] = -1; // initialize dimension counter + for (int64_t dim_idx = 0; dim_idx < n_dim; dim_idx++) { + input_idx += input_mappings[dim_idx][0 /* output_dim_counter[dim_idx] */]; + } - for (; output_idx < output_shape.Size(); output_idx++) { + for (int64_t output_size = output_shape.Size(); output_idx < output_size; output_idx++) { + output[output_idx] = (input_idx < 0) ? static_cast(extrapolation_value) : input[input_idx]; for (int64_t dim_idx = n_dim - 1; dim_idx >= 0; dim_idx--) { + input_idx -= input_mappings[dim_idx][output_dim_counter[dim_idx]]; if (++output_dim_counter[dim_idx] < output_shape[dim_idx]) { - int64_t current_input_dim_counter = 0; - auto original_idx = get_original_coordinate(static_cast(output_dim_counter[dim_idx]), scales[dim_idx], - static_cast(output_shape[dim_idx]), static_cast(input_shape[dim_idx]), - roi[dim_idx], roi[n_dim + dim_idx]); - current_input_dim_counter = get_nearest_pixel(original_idx, scales[dim_idx] < 1); - current_input_dim_counter = std::max((int64_t)0, - std::min(current_input_dim_counter, (input_shape[dim_idx] - 1))); - - if (current_input_dim_counter != input_dim_counters[dim_idx]) { - input_idx += (current_input_dim_counter - input_dim_counters[dim_idx]) * input_dim_factor[dim_idx]; - input_dim_counters[dim_idx] = current_input_dim_counter; - } + input_idx += input_mappings[dim_idx][output_dim_counter[dim_idx]]; break; - } else { - output_dim_counter[dim_idx] = 0; - input_idx += (0 - input_dim_counters[dim_idx]) * input_dim_factor[dim_idx]; - input_dim_counters[dim_idx] = 0; } + output_dim_counter[dim_idx] = 0; + input_idx += input_mappings[dim_idx][0 /* output_dim_counter[dim_idx] */ ]; } - - output[output_idx] = input[input_idx]; } return Status::OK(); diff --git a/onnxruntime/test/providers/cpu/tensor/resize_op_test.cc b/onnxruntime/test/providers/cpu/tensor/resize_op_test.cc index 7f8c9b8c6d..72f7b03074 100644 --- a/onnxruntime/test/providers/cpu/tensor/resize_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/resize_op_test.cc @@ -327,6 +327,35 @@ TEST(ResizeOpTest, ResizeOpNearestDownSampleTest_tf_crop_and_resize_with_extrapo test.Run(); } +TEST(ResizeOpTest, ResizeOpNearestDownSample5dTest_tf_crop_and_resize_with_extrapolation) { + OpTester test("Resize", 11); + std::vector scales{1.0f, 1.0f, 1.0f, 0.8f, 0.8f}; + std::vector roi{0.0f, 0.0f, 0.0f, 0.4f, 0.6f, 1.0f, 1.0f, 1.0f, 1.2f, 1.7f}; + + test.AddAttribute("mode", "nearest"); + test.AddAttribute("coordinate_transformation_mode", "tf_crop_and_resize"); + test.AddAttribute("extrapolation_value", 10.0f); + + const int64_t N = 1, C = 1, H = 4, W = 4; + std::vector X = { + 1.0f, 2.0f, 3.0f, 4.0f, + 5.0f, 6.0f, 7.0f, 8.0f, + 9.0f, 10.0f, 11.0f, 12.0f, + 13.0f, 14.0f, 15.0f, 16.0f}; + + test.AddInput("X", {1, N, C, H, W}, X); + test.AddInput("roi", {10}, roi); + test.AddInput("scales", {5}, scales); + + std::vector Y = {7.0f, 10.0f, 10.0f, + 11.0f, 10.f, 10.0f, + 10.0f, 10.0f, 10.0f}; + + test.AddOutput("Y", {1, N, C, static_cast(H * scales[3]), static_cast(W * scales[4])}, Y); + // Current cuda provider do not support more than 4d + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCudaExecutionProvider}); +} + TEST(ResizeOpTest, ResizeOpNearestUpSampleTest) { OpTester test("Resize", 11); std::vector roi{}; @@ -379,6 +408,37 @@ TEST(ResizeOpTest, ResizeOpNearestUpSampleTest_WithSizes_CeilMode) { test.Run(); } +TEST(ResizeOpTest, ResizeOpNearestUpSample5dTest_WithSizes_CeilMode) { + OpTester test("Resize", 11); + std::vector roi{}; + std::vector scales{}; + std::vector sizes{1, 1, 1, 7, 8}; + + test.AddAttribute("mode", "nearest"); + test.AddAttribute("nearest_mode", "ceil"); + + const int64_t N = 1, C = 1, H = 2, W = 2; + std::vector X = {1.0f, 2.0f, 3.0f, 4.0f}; + + test.AddInput("X", {1, N, C, H, W}, X); + test.AddInput("roi", {0}, roi); + test.AddInput("scales", {0}, scales); + test.AddInput("sizes", {5}, sizes); + + std::vector Y = {1.0f, 1.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 1.0f, 1.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, 2.0f, + 3.0f, 3.0f, 4.0f, 4.0f, 4.0f, 4.0f, 4.0f, 4.0f, + 3.0f, 3.0f, 4.0f, 4.0f, 4.0f, 4.0f, 4.0f, 4.0f, + 3.0f, 3.0f, 4.0f, 4.0f, 4.0f, 4.0f, 4.0f, 4.0f, + 3.0f, 3.0f, 4.0f, 4.0f, 4.0f, 4.0f, 4.0f, 4.0f, + 3.0f, 3.0f, 4.0f, 4.0f, 4.0f, 4.0f, 4.0f, 4.0f}; + + test.AddOutput("Y", {1, N, C, sizes[3], sizes[4]}, Y); + // Current cuda provider do not support more than 4d + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kCudaExecutionProvider}); +} + + TEST(ResizeOpTest, ResizeOpNearestUpSample_Floor_Align_Corners) { OpTester test("Resize", 11);