diff --git a/onnxruntime/core/providers/cpu/tensor/upsample.cc b/onnxruntime/core/providers/cpu/tensor/upsample.cc index e5b6dc938f..f6fa2195f8 100644 --- a/onnxruntime/core/providers/cpu/tensor/upsample.cc +++ b/onnxruntime/core/providers/cpu/tensor/upsample.cc @@ -22,13 +22,14 @@ ONNX_CPU_OPERATOR_VERSIONED_TYPED_KERNEL( KernelDefBuilder().TypeConstraint("T", DataTypeImpl::GetTensorType()), Upsample); -void upsampleNearest2x( +template +void UpsampleNearest2x( int64_t batch_size, int64_t num_channels, int64_t input_height, int64_t input_width, - const float* input, - float* output) { + const T* input, + T* output) { const int64_t output_height = input_height * 2; const int64_t output_width = input_width * 2; for (int64_t n = 0; n < batch_size; ++n) { @@ -36,7 +37,7 @@ void upsampleNearest2x( for (int64_t y = 0; y < output_height; ++y) { const int64_t in_y = y / 2; for (int64_t x = 0; x < input_width; ++x) { - const float v = input[in_y * input_width + x]; + const T v = input[in_y * input_width + x]; const int64_t oidx = output_width * y + x * 2; output[oidx + 0] = v; output[oidx + 1] = v; @@ -49,7 +50,7 @@ void upsampleNearest2x( } template -Status upsampleNearest(const T* input, +Status UpsampleNearest(const T* input, T* output, const TensorShape& input_shape, const TensorShape& output_shape, @@ -59,19 +60,23 @@ Status upsampleNearest(const T* input, if (input_shape.NumDimensions() != output_shape.NumDimensions()) return Status(ONNXRUNTIME, FAIL, "Upsample: input/output value's dimension mismatch"); auto n_dim = input_shape.NumDimensions(); - for (size_t i = 0, size = output_shape.Size(); i < size; i++) { - size_t old_idx = 0; - size_t cur_idx = i; + if (scales.size() == 4 && scales[0] == 1 && scales[1] == 1 && scales[2] == 2 && scales[3] == 2) { + UpsampleNearest2x(input_shape[0], input_shape[1], input_shape[2], input_shape[3], input, output); + } else { + for (size_t i = 0, size = output_shape.Size(); i < size; i++) { + size_t old_idx = 0; + size_t cur_idx = i; - int64_t base = 1; - for (int64_t j = static_cast(n_dim - 1); j >= 0; j--) { - auto tmp = cur_idx % output_shape[j]; - old_idx += (std::min(static_cast(tmp / scales[j]), input_shape[j] - 1)) * base; - base *= input_shape[j]; - cur_idx /= output_shape[j]; + int64_t base = 1; + for (int64_t j = static_cast(n_dim - 1); j >= 0; j--) { + auto tmp = cur_idx % output_shape[j]; + old_idx += (std::min(static_cast(tmp / scales[j]), input_shape[j] - 1)) * base; + base *= input_shape[j]; + cur_idx /= output_shape[j]; + } + + output[i] = input[old_idx]; } - - output[i] = input[old_idx]; } return Status::OK(); } @@ -208,7 +213,7 @@ Status Upsample::BaseCompute(OpKernelContext* context, const std::vector(X->template Data(), Y->template MutableData(), X->Shape(), Y->Shape(), scales); + return UpsampleNearest(X->template Data(), Y->template MutableData(), X->Shape(), Y->Shape(), scales); case UpsampleMode::LINEAR: { //What's the correct behavior of linear mode is not clear right now, //Only support bilinear with 4D tensor to keep consistent with previous behavior @@ -227,7 +232,6 @@ Status Upsample::BaseCompute(OpKernelContext* context, const std::vector Status Upsample::Compute(OpKernelContext* context) const { if (OpKernel::Node().InputDefs().size() == 1 || scales_cached_) { diff --git a/onnxruntime/core/providers/cpu/tensor/upsample.h b/onnxruntime/core/providers/cpu/tensor/upsample.h index 2f65b22947..957c254147 100644 --- a/onnxruntime/core/providers/cpu/tensor/upsample.h +++ b/onnxruntime/core/providers/cpu/tensor/upsample.h @@ -17,7 +17,7 @@ enum UpsampleMode { class UpsampleBase { protected: - UpsampleBase(OpKernelInfo info): scales_cached_(false) { + UpsampleBase(OpKernelInfo info) : scales_cached_(false) { std::string mode; ORT_ENFORCE(info.GetAttr("mode", &mode).IsOK()); mode_ = StringToUpsampleMode(mode); @@ -51,7 +51,7 @@ class UpsampleBase { return UpsampleMode::LINEAR; } else { ORT_THROW("mode attribute is " + mode + ". It can only be " + - UpsampleModeNN + "(default) or " + UpsampleModeLinear + "."); + UpsampleModeNN + "(default) or " + UpsampleModeLinear + "."); } } @@ -63,7 +63,7 @@ class UpsampleBase { if (UpsampleMode::LINEAR == mode) { ORT_ENFORCE(scales.size() == 4, "Upsample: linear mode upsample only support bilinear with 4 dimension."); ORT_ENFORCE(((scales[0] == 1) && (scales[1] == 1)), - "Upsample: linear mode upsample only support bilinear, the first 2 scales should be 1."); + "Upsample: linear mode upsample only support bilinear, the first 2 scales should be 1."); } } @@ -71,6 +71,9 @@ class UpsampleBase { const float* scale_data = scale->template Data(); int64_t scales_size = scale->Shape().Size(); ORT_ENFORCE(scales_size > 0, "scales size should be greater than 0."); + if (scales.size() == 0) { + scales.resize(scales_size); + } memcpy(scales.data(), scale_data, scales_size * sizeof(float)); ScalesValidation(scales, mode_); }