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

implement gridsample 20 (#17744)

Browse source
This commit is contained in:
liqun Fu 2023-11-07 10:42:41 -08:00 committed by GitHub
parent 83c0275354
commit 6127dd1d2d
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
9 changed files with 1379 additions and 154 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
docs/OperatorKernels.md
View file

@ -138,7 +138,8 @@ Do not modify directly.*
|||[7, 8]|**T** = tensor(double), tensor(float)<br/> **T1** = tensor(bool)|
|GreaterOrEqual|*in* A:**T**<br> *in* B:**T**<br> *out* C:**T1**|16+|**T** = tensor(double), tensor(float), tensor(int32), tensor(int64)<br/> **T1** = tensor(bool)|
|||[12, 15]|**T** = tensor(double), tensor(float), tensor(int32), tensor(int64)<br/> **T1** = tensor(bool)|
|GridSample|*in* X:**T1**<br> *in* grid:**T2**<br> *out* Y:**T1**|16+|**T1** = tensor(float)<br/> **T2** = tensor(float)|
|GridSample|*in* X:**T1**<br> *in* grid:**T2**<br> *out* Y:**T1**|20+|**T1** = tensor(double), tensor(float)<br/> **T2** = tensor(double), tensor(float)|
|||[16, 19]|**T1** = tensor(float)<br/> **T2** = tensor(float)|
|HammingWindow|*in* size:**T1**<br> *out* output:**T2**|17+|**T1** = tensor(int32), tensor(int64)<br/> **T2** = tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|HannWindow|*in* size:**T1**<br> *out* output:**T2**|17+|**T1** = tensor(int32), tensor(int64)<br/> **T2** = tensor(double), tensor(float), tensor(int16), tensor(int32), tensor(int64), tensor(int8), tensor(uint16), tensor(uint32), tensor(uint64), tensor(uint8)|
|HardSigmoid|*in* X:**T**<br> *out* Y:**T**|6+|**T** = tensor(float)|

10
onnxruntime/core/providers/cpu/cpu_execution_provider.cc
View file

@ -798,7 +798,7 @@ class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDoma
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, 18, If);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, float, RoiAlign);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, double, RoiAlign);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, float, GridSample);
class ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, 19, float, GridSample);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, 17, ScatterElements);
class ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, 17, ScatterND);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, string, Where);
@ -960,6 +960,8 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 19, Sh
// Opset 20
class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, ConstantOfShape);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, float, GridSample);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, double, GridSample);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, float, AffineGrid);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, double, AffineGrid);
class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, float, IsNaN);
@ -2183,8 +2185,8 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
RoiAlign)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, double,
RoiAlign)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, float,
GridSample)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, 19, float,
GridSample)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, 17, ScatterElements)>,
BuildKernelCreateInfo<ONNX_OPERATOR_VERSIONED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, 17, ScatterND)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 16, string, Where)>,
@ -2401,6 +2403,8 @@ Status RegisterOnnxOperatorKernels(KernelRegistry& kernel_registry) {
// Opset 20
BuildKernelCreateInfo<ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, ConstantOfShape)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, float, GridSample)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, double, GridSample)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, float, AffineGrid)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, double, AffineGrid)>,
BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 20, float, IsNaN)>,

338
onnxruntime/core/providers/cpu/tensor/grid_sample.cc
View file

@ -11,17 +11,23 @@
namespace onnxruntime {
#define REGISTER_KERNEL_TYPED(T) \
ONNX_CPU_OPERATOR_TYPED_KERNEL( \
GridSample, \
16, \
T, \
KernelDefBuilder() \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<T>()), \
GridSample<T>);
#define REGISTER_KERNEL_TYPED(T) \
ONNX_OPERATOR_VERSIONED_TYPED_KERNEL_EX(GridSample, kOnnxDomain, 16, 19, T, kCpuExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<T>()), \
GridSample<T>);
#define REGISTER_KERNEL_TYPED_20(T) \
ONNX_OPERATOR_TYPED_KERNEL_EX(GridSample, kOnnxDomain, 20, T, kCpuExecutionProvider, \
KernelDefBuilder() \
.TypeConstraint("T1", DataTypeImpl::GetTensorType<T>()) \
.TypeConstraint("T2", DataTypeImpl::GetTensorType<T>()), \
GridSample<T>);
REGISTER_KERNEL_TYPED(float)
REGISTER_KERNEL_TYPED_20(float)
REGISTER_KERNEL_TYPED_20(double)
// Restore normalized location to actual image location
// When align_corners is true:
@ -44,16 +50,15 @@ T GsDenormalize(T n, int64_t length, bool align_corners) {
}
// Reflect by the near border till within the borders
// Use float for borders to avoid potential issues with integer T
template <typename T>
T GsReflect(T x, float x_min, float x_max) {
float dx = {};
float fx = static_cast<float>(x);
float range = x_max - x_min;
T GsReflect(T x, T x_min, T x_max) {
T dx = {};
T fx = static_cast<T>(x);
T range = x_max - x_min;
if (fx < x_min) {
dx = x_min - fx;
int n = static_cast<int>(dx / range);
float r = dx - n * range;
T r = dx - n * range;
if (n % 2 == 0) {
fx = x_min + r;
} else {
@ -62,7 +67,7 @@ T GsReflect(T x, float x_min, float x_max) {
} else if (fx > x_max) {
dx = fx - x_max;
int n = static_cast<int>(dx / range);
float r = dx - n * range;
T r = dx - n * range;
if (n % 2 == 0) {
fx = x_max - r;
} else {
@ -75,9 +80,9 @@ T GsReflect(T x, float x_min, float x_max) {
// Calculate cubic convolution interpolation coefficients
// ROBERT G. KEYS https://ieeexplore.ieee.org/document/1163711
// Use float to avoid potential issues with integer T
void GsGetCubicCoeffs(float x, float coeffs[4]) {
constexpr float cubic_alpha = -0.75f;
template <typename T>
void GsGetCubicCoeffs(T x, T coeffs[4]) {
constexpr T cubic_alpha = -0.75f;
x = std::abs(x);
coeffs[0] = ((cubic_alpha * (x + 1) - 5 * cubic_alpha) * (x + 1) + 8 * cubic_alpha) * (x + 1) - 4 * cubic_alpha;
coeffs[1] = ((cubic_alpha + 2) * x - (cubic_alpha + 3)) * x * x + 1;
@ -86,9 +91,9 @@ void GsGetCubicCoeffs(float x, float coeffs[4]) {
}
template <typename T>
T GsBicubicInterpolate(T p[4][4], float x, float y) {
float v[4] = {};
float coeffs[4] = {};
T GsBicubicInterpolate(T p[4][4], T x, T y) {
T v[4] = {};
T coeffs[4] = {};
GsGetCubicCoeffs(x, coeffs);
for (int64_t i = 0; i < 4; i++) {
v[i] = coeffs[0] * p[i][0] + coeffs[1] * p[i][1] + coeffs[2] * p[i][2] + coeffs[3] * p[i][3];
@ -98,7 +103,7 @@ T GsBicubicInterpolate(T p[4][4], float x, float y) {
}
template <typename T>
T GridSample<T>::PixelAtGrid(const T* image, int64_t r, int64_t c, int64_t H, int64_t W, float border[/* 4 */]) const {
T GridSample<T>::PixelAtGrid(const T* image, int64_t r, int64_t c, int64_t H, int64_t W, T border[/* 4 */]) const {
T pixel = {}; // default 0
if (padding_mode_ == Zeros) {
if (c >= 0 && c < W && r >= 0 && r < H) {
@ -116,6 +121,27 @@ T GridSample<T>::PixelAtGrid(const T* image, int64_t r, int64_t c, int64_t H, in
return pixel;
}
template <typename T>
T GridSample<T>::PixelAtGrid3D(const T* image, int64_t d, int64_t h, int64_t w, int64_t D, int64_t H, int64_t W, T border[/* 6 */]) const {
T pixel = {}; // default 0
if (padding_mode_ == Zeros) {
if (w >= 0 && w < W && h >= 0 && h < H && d >= 0 && d < D) {
pixel = image[d * H * W + h * W + w];
}
} else if (padding_mode_ == Border) {
w = std::clamp<int64_t>(w, 0, W - 1);
h = std::clamp<int64_t>(h, 0, H - 1);
d = std::clamp<int64_t>(d, 0, D - 1);
pixel = image[d * H * W + h * W + w];
} else { // (padding_mode_ == Reflection)
w = static_cast<int64_t>(GsReflect(static_cast<T>(w), border[0], border[3]));
h = static_cast<int64_t>(GsReflect(static_cast<T>(h), border[1], border[4]));
d = static_cast<int64_t>(GsReflect(static_cast<T>(d), border[2], border[5]));
pixel = image[d * H * W + h * W + w];
}
return pixel;
}
// When grid sampling, padding is applied before interpolation.
// For instance, in bilinear mode and zeros padding-mode, pixel p at actual
// image location (-0.5, -0.5)
@ -134,113 +160,203 @@ Status GridSample<T>::Compute(OpKernelContext* context) const {
const auto& input_dims = input->Shape();
const auto& grid_dims = grid->Shape();
if (input_dims.NumDimensions() != 4 || grid_dims.NumDimensions() != 4) {
return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "Only 4-D tensor is supported");
}
int64_t data_dims = input_dims.NumDimensions() - 2;
ORT_ENFORCE(static_cast<int64_t>(grid_dims.NumDimensions()) == data_dims + 2,
"grid dimensions must be ", data_dims + 2, "for input dimension of ", data_dims);
ORT_ENFORCE(grid_dims[grid_dims.NumDimensions() - 1] == data_dims,
"Last dimension of grid: ", grid_dims[grid_dims.NumDimensions() - 1], ", expect ", data_dims);
ORT_ENFORCE(input_dims.NumDimensions() == 4 || input_dims.NumDimensions() == 5, "Only 4-D or 5-D tensor is supported");
auto N = input_dims[0];
auto C = input_dims[1];
auto H_in = input_dims[2];
auto W_in = input_dims[3];
auto H_out = grid_dims[1];
auto W_out = grid_dims[2];
ORT_ENFORCE(grid_dims[0] == N, "Grid batch size ", grid_dims[0], " does not match input batch size ", N);
ORT_ENFORCE(grid_dims[3] == 2, "Last dimension of grid: ", grid_dims[3], ", expect 2");
TensorShape Y_shape = {N, C, H_out, W_out};
auto& Y = *context->Output(0, Y_shape);
// Return early if the output tensor is going to be of size 0
if (Y.Shape().Size() == 0) {
return Status::OK();
if (input_dims.NumDimensions() == 5) {
ORT_ENFORCE(mode_ != Cubic, "Only support GridSample Cubic mode in 4-D cases.");
}
// Force float here to avoid possible issue in integer T case
float x_min = -0.5f;
float x_max = W_in - 0.5f;
float y_min = -0.5f;
float y_max = H_in - 0.5f;
if (data_dims == 2) {
// sample 2d;
auto H_in = input_dims[2];
auto W_in = input_dims[3];
auto H_out = grid_dims[1];
auto W_out = grid_dims[2];
TensorShape Y_shape = {N, C, H_out, W_out};
auto& Y = *context->Output(0, Y_shape);
// Return early if the output tensor is going to be of size 0
if (Y.Shape().Size() == 0) {
return Status::OK();
}
if (align_corners_) {
x_min = 0.f;
x_max = W_in - 1.f;
y_min = 0.f;
y_max = H_in - 1.f;
}
float border[] = {x_min, y_min, x_max, y_max}; // l-t-r-b
T x_min = -0.5f;
T x_max = W_in - 0.5f;
T y_min = -0.5f;
T y_max = H_in - 0.5f;
concurrency::ThreadPool* tp = H_out * W_out > 64 ? context->GetOperatorThreadPool() : nullptr;
for (int64_t n = 0; n < N; n++) {
const T* grid_data = grid->Data<T>() + n * (H_out * W_out) * 2;
concurrency::ThreadPool::TrySimpleParallelFor(
tp, onnxruntime::narrow<std::ptrdiff_t>(C),
[&](std::ptrdiff_t c) {
const T* X_data = input->Data<T>() + (n * C + c) * (H_in * W_in);
T* Y_data = Y.MutableData<T>() + (n * C + c) * (H_out * W_out);
if (align_corners_) {
x_min = 0.f;
x_max = W_in - 1.f;
y_min = 0.f;
y_max = H_in - 1.f;
}
T border[] = {x_min, y_min, x_max, y_max}; // l-t-r-b
for (int64_t oy = 0; oy < H_out; oy++) {
for (int64_t ox = 0; ox < W_out; ox++) {
const T* gridpoint = grid_data + (oy * W_out + ox) * 2;
T* Y_gridpoint = Y_data + oy * W_out + ox;
auto nx = gridpoint[0]; // normalized location
auto ny = gridpoint[1];
auto x = GsDenormalize<T>(nx, W_in, align_corners_); // actual location
auto y = GsDenormalize<T>(ny, H_in, align_corners_);
concurrency::ThreadPool* tp = H_out * W_out > 64 ? context->GetOperatorThreadPool() : nullptr;
for (int64_t n = 0; n < N; n++) {
const T* grid_data = grid->Data<T>() + n * (H_out * W_out) * 2;
concurrency::ThreadPool::TrySimpleParallelFor(
tp, onnxruntime::narrow<std::ptrdiff_t>(C),
[&](std::ptrdiff_t c) {
const T* X_data = input->Data<T>() + (n * C + c) * (H_in * W_in);
T* Y_data = Y.MutableData<T>() + (n * C + c) * (H_out * W_out);
if (mode_ == Nearest) {
x = static_cast<T>(std::nearbyintf(static_cast<float>(x)));
y = static_cast<T>(std::nearbyintf(static_cast<float>(y)));
}
for (int64_t oy = 0; oy < H_out; oy++) {
for (int64_t ox = 0; ox < W_out; ox++) {
const T* gridpoint = grid_data + (oy * W_out + ox) * 2;
T* Y_gridpoint = Y_data + oy * W_out + ox;
auto nx = gridpoint[0]; // normalized location
auto ny = gridpoint[1];
auto x = GsDenormalize<T>(nx, W_in, align_corners_); // actual location
auto y = GsDenormalize<T>(ny, H_in, align_corners_);
if (x < x_min || x > x_max || y < y_min || y > y_max) { // out of bound
if (padding_mode_ == Border) {
// use original border in both align_corner cases
x = std::clamp(x, static_cast<T>(0), static_cast<T>(W_in - 1));
y = std::clamp(y, static_cast<T>(0), static_cast<T>(H_in - 1));
} else if (padding_mode_ == Reflection) {
x = GsReflect(x, x_min, x_max);
y = GsReflect(y, y_min, y_max);
}
} // out of bound
if (mode_ == Nearest) {
x = static_cast<T>(std::nearbyint(static_cast<T>(x)));
y = static_cast<T>(std::nearbyint(static_cast<T>(y)));
// x, y are integers in all padding modes
*Y_gridpoint = PixelAtGrid(X_data, static_cast<int64_t>(y), static_cast<int64_t>(x), H_in, W_in, border);
} else if (mode_ == Linear) {
int64_t x1 = static_cast<int64_t>(std::floor(x));
int64_t y1 = static_cast<int64_t>(std::floor(y));
int64_t x2 = x1 + 1;
int64_t y2 = y1 + 1;
if (mode_ == Nearest) {
// x, y are integers in all padding modes
*Y_gridpoint = PixelAtGrid(X_data, static_cast<int64_t>(y), static_cast<int64_t>(x), H_in, W_in, border);
continue;
}
T p11 = PixelAtGrid(X_data, y1, x1, H_in, W_in, border);
T p12 = PixelAtGrid(X_data, y1, x2, H_in, W_in, border);
T p21 = PixelAtGrid(X_data, y2, x1, H_in, W_in, border);
T p22 = PixelAtGrid(X_data, y2, x2, H_in, W_in, border);
if (mode_ == Bilinear) {
int64_t x1 = static_cast<int64_t>(std::floor(x));
int64_t y1 = static_cast<int64_t>(std::floor(y));
int64_t x2 = x1 + 1;
int64_t y2 = y1 + 1;
T dx2 = static_cast<T>(x2) - x;
T dx1 = x - static_cast<T>(x1);
T dy2 = static_cast<T>(y2) - y;
T dy1 = y - static_cast<T>(y1);
*Y_gridpoint = dy2 * (dx2 * p11 + dx1 * p12) + dy1 * (dx2 * p21 + dx1 * p22);
} else if (mode_ == Cubic) {
int64_t x0 = static_cast<int64_t>(std::floor(x)) - 1; // top-left corner of the bbox
int64_t y0 = static_cast<int64_t>(std::floor(y)) - 1;
T p11 = PixelAtGrid(X_data, y1, x1, H_in, W_in, border);
T p12 = PixelAtGrid(X_data, y1, x2, H_in, W_in, border);
T p21 = PixelAtGrid(X_data, y2, x1, H_in, W_in, border);
T p22 = PixelAtGrid(X_data, y2, x2, H_in, W_in, border);
T dx2 = static_cast<T>(x2) - x;
T dx1 = x - static_cast<T>(x1);
T dy2 = static_cast<T>(y2) - y;
T dy1 = y - static_cast<T>(y1);
*Y_gridpoint = dy2 * (dx2 * p11 + dx1 * p12) + dy1 * (dx2 * p21 + dx1 * p22);
}
if (mode_ == Bicubic) {
int64_t x0 = static_cast<int64_t>(std::floor(x)) - 1; // top-left corner of the bbox
int64_t y0 = static_cast<int64_t>(std::floor(y)) - 1;
T p[4][4] = {}; // [H][W]
for (int64_t h = 0; h < 4; h++) {
for (int64_t w = 0; w < 4; w++) {
p[h][w] = PixelAtGrid(X_data, h + y0, w + x0, H_in, W_in, border);
T p[4][4] = {}; // [H][W]
for (int64_t h = 0; h < 4; h++) {
for (int64_t w = 0; w < 4; w++) {
p[h][w] = PixelAtGrid(X_data, h + y0, w + x0, H_in, W_in, border);
}
}
T dx = static_cast<T>(x - x0 - 1);
T dy = static_cast<T>(y - y0 - 1);
*Y_gridpoint = GsBicubicInterpolate(p, dx, dy);
}
T dx = static_cast<T>(x - x0 - 1);
T dy = static_cast<T>(y - y0 - 1);
*Y_gridpoint = GsBicubicInterpolate(p, static_cast<float>(dx), static_cast<float>(dy));
}
}
}
});
});
}
} else if (data_dims == 3) {
// sample 3d;
auto D_in = input_dims[2];
auto H_in = input_dims[3];
auto W_in = input_dims[4];
auto D_out = grid_dims[1];
auto H_out = grid_dims[2];
auto W_out = grid_dims[3];
TensorShape Y_shape = {N, C, D_out, H_out, W_out};
auto& Y = *context->Output(0, Y_shape);
// Return early if the output tensor is going to be of size 0
if (Y.Shape().Size() == 0) {
return Status::OK();
}
T x_min = -0.5f;
T x_max = W_in - 0.5f;
T y_min = -0.5f;
T y_max = H_in - 0.5f;
T z_min = -0.5f;
T z_max = D_in - 0.5f;
if (align_corners_) {
x_min = 0.f;
x_max = W_in - 1.f;
y_min = 0.f;
y_max = H_in - 1.f;
z_min = 0.f;
z_max = D_in - 1.f;
}
T border[] = {x_min, y_min, z_min, x_max, y_max, z_max};
concurrency::ThreadPool* tp = D_out * H_out * W_out > 64 ? context->GetOperatorThreadPool() : nullptr;
for (int64_t n = 0; n < N; n++) {
const T* grid_data = grid->Data<T>() + n * (D_out * H_out * W_out) * 3;
concurrency::ThreadPool::TrySimpleParallelFor(
tp, onnxruntime::narrow<std::ptrdiff_t>(C),
[&](std::ptrdiff_t c) {
const T* X_data = input->Data<T>() + (n * C + c) * (D_in * H_in * W_in);
T* Y_data = Y.MutableData<T>() + (n * C + c) * (D_out * H_out * W_out);
for (int64_t oz = 0; oz < D_out; oz++) {
for (int64_t oy = 0; oy < H_out; oy++) {
for (int64_t ox = 0; ox < W_out; ox++) {
const T* gridpoint = grid_data + (oz * H_out * W_out + oy * W_out + ox) * 3;
T* Y_gridpoint = Y_data + oz * H_out * W_out + oy * W_out + ox;
auto nx = gridpoint[0]; // normalized location
auto ny = gridpoint[1];
auto nz = gridpoint[2];
auto x = GsDenormalize<T>(nx, W_in, align_corners_); // actual location
auto y = GsDenormalize<T>(ny, H_in, align_corners_);
auto z = GsDenormalize<T>(nz, D_in, align_corners_);
if (mode_ == Nearest) {
x = static_cast<T>(std::nearbyint(static_cast<T>(x)));
y = static_cast<T>(std::nearbyint(static_cast<T>(y)));
z = static_cast<T>(std::nearbyint(static_cast<T>(z)));
// x, y are integers in all padding modes
*Y_gridpoint = PixelAtGrid3D(X_data, static_cast<int64_t>(z), static_cast<int64_t>(y), static_cast<int64_t>(x),
D_in, H_in, W_in, border);
} else if (mode_ == Linear) {
int64_t x1 = static_cast<int64_t>(std::floor(x));
int64_t y1 = static_cast<int64_t>(std::floor(y));
int64_t z1 = static_cast<int64_t>(std::floor(z));
int64_t x2 = x1 + 1;
int64_t y2 = y1 + 1;
int64_t z2 = z1 + 1;
T dx2 = static_cast<T>(x2) - x;
T dx1 = x - static_cast<T>(x1);
T dy2 = static_cast<T>(y2) - y;
T dy1 = y - static_cast<T>(y1);
T dz2 = static_cast<T>(z2) - z;
T dz1 = z - static_cast<T>(z1);
T p111 = PixelAtGrid3D(X_data, z1, y1, x1, D_in, H_in, W_in, border);
T p112 = PixelAtGrid3D(X_data, z1, y1, x2, D_in, H_in, W_in, border);
T p121 = PixelAtGrid3D(X_data, z1, y2, x1, D_in, H_in, W_in, border);
T p122 = PixelAtGrid3D(X_data, z1, y2, x2, D_in, H_in, W_in, border);
T Y_gridpoint_z1 = dy2 * (dx2 * p111 + dx1 * p112) + dy1 * (dx2 * p121 + dx1 * p122);
T p211 = PixelAtGrid3D(X_data, z2, y1, x1, D_in, H_in, W_in, border);
T p212 = PixelAtGrid3D(X_data, z2, y1, x2, D_in, H_in, W_in, border);
T p221 = PixelAtGrid3D(X_data, z2, y2, x1, D_in, H_in, W_in, border);
T p222 = PixelAtGrid3D(X_data, z2, y2, x2, D_in, H_in, W_in, border);
T Y_gridpoint_z2 = dy2 * (dx2 * p211 + dx1 * p212) + dy1 * (dx2 * p221 + dx1 * p222);
*Y_gridpoint = dz2 * Y_gridpoint_z1 + dz1 * Y_gridpoint_z2;
}
}
}
}
});
}
} else {
// shall not reach here due to above checks
return Status(common::ONNXRUNTIME, common::INVALID_ARGUMENT, "Only support GirdSample in 4-D or 5-D cases.");
}
return Status::OK();
}

54
onnxruntime/core/providers/cpu/tensor/grid_sample.h
View file

@ -15,37 +15,52 @@ template <typename T>
class GridSample final : public OpKernel {
public:
explicit GridSample(const OpKernelInfo& info) : OpKernel(info) {
std::string mode_str = info.GetAttrOrDefault<std::string>("mode", "bilinear");
int start_version = info.node().SinceVersion();
if (start_version >= 20) {
std::string mode_str = info.GetAttrOrDefault<std::string>("mode", "linear");
if (mode_str == "cubic") {
mode_ = Cubic;
} else if (mode_str == "nearest") {
mode_ = Nearest;
} else if (mode_str == "linear") {
mode_ = Linear;
} else {
ORT_THROW("mode \"", mode_str, "\" not supported, expect linear, nearest or cubic");
}
} else {
std::string mode_str = info.GetAttrOrDefault<std::string>("mode", "bilinear");
if (mode_str == "bicubic") {
mode_ = Cubic;
} else if (mode_str == "nearest") {
mode_ = Nearest;
} else if (mode_str == "bilinear") {
mode_ = Linear;
} else {
ORT_THROW("mode \"", mode_str, "\" not supported, expect bilinear, nearest or bicubic");
}
}
std::string padding_mode_str = info.GetAttrOrDefault<std::string>("padding_mode", "zeros");
align_corners_ = static_cast<bool>(info.GetAttrOrDefault<int64_t>("align_corners", 0));
ORT_ENFORCE(mode_str == "bilinear" || mode_str == "nearest" || mode_str == "bicubic",
"mode \"", mode_str, "\" not supported, expect bilinear, nearest or bicubic");
ORT_ENFORCE(padding_mode_str == "zeros" || padding_mode_str == "border" || padding_mode_str == "reflection",
"padding_mode \"", padding_mode_str, "\" not supported, expect zeros, border or reflection");
if (mode_str == "bicubic") {
mode_ = Bicubic;
} else if (mode_str == "nearest") {
mode_ = Nearest;
} else {
mode_ = Bilinear;
}
if (padding_mode_str == "reflection") {
padding_mode_ = Reflection;
} else if (padding_mode_str == "border") {
padding_mode_ = Border;
} else {
} else if (padding_mode_str == "zeros") {
padding_mode_ = Zeros;
} else {
ORT_THROW("padding_mode \"", padding_mode_str, "\" not supported, expect zeros, border or reflection");
}
}
Status Compute(OpKernelContext* context) const override;
private:
enum GridSampleInterpolationMode {
Bilinear,
typedef enum {
Linear,
Cubic,
Nearest,
Bicubic
};
} GridSampleInterpolationMode;
enum GridSamplePaddingMode {
Zeros,
@ -53,9 +68,10 @@ class GridSample final : public OpKernel {
Reflection
};
T PixelAtGrid(const T* image, int64_t r, int64_t c, int64_t H, int64_t W, float border[/* 4 */]) const;
T PixelAtGrid(const T* image, int64_t r, int64_t c, int64_t H, int64_t W, T border[/* 4 */]) const;
T PixelAtGrid3D(const T* image, int64_t d, int64_t h, int64_t w, int64_t D, int64_t H, int64_t W, T border[/* 6 */]) const;
GridSampleInterpolationMode mode_{Bilinear};
GridSampleInterpolationMode mode_{Linear};
GridSamplePaddingMode padding_mode_{Zeros};
bool align_corners_{0};
};

4
onnxruntime/test/onnx/TestCase.cc
View file

@ -944,8 +944,6 @@ std::unique_ptr<std::set<BrokenTest>> GetBrokenTests(const std::string& provider
{"simple_rnn_batchwise", "type error", {}},
{"mod_float_mixed_sign_example", "fmod attribute must be true for floating point types", {}},
{"col2im_pads", "result mismatch", {"opset18"}},
{"gridsample_volumetric_nearest_align_corners_0", "result differs", {}},
{"gridsample_volumetric_nearest_align_corners_1", "result differs", {}},
{"reduce_l1_empty_set", "unknown version", {}},
{"reduce_l1_empty_set_expanded", "unknown version", {}},
{"reduce_l2_empty_set", "unknown version", {}},
@ -1351,6 +1349,8 @@ std::unique_ptr<std::set<BrokenTest>> GetBrokenTests(const std::string& provider
broken_tests->insert({"sce_sum_log_prob", "result differs"});
broken_tests->insert({"sce_sum_log_prob_expanded", "result differs"});
broken_tests->insert({"gridsample_reflection_padding", "result differs"});
broken_tests->insert({"gridsample_volumetric_nearest_align_corners_0", "unknown version"});
broken_tests->insert({"gridsample_volumetric_nearest_align_corners_1", "unknown version"});
broken_tests->insert({"spacetodepth", "result differs"});
}

6
onnxruntime/test/providers/cpu/tensor/affine_grid_test_gen.py
View file

@ -1,3 +1,9 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# This code is used to generate the test cases for the AffineGrid operator
# in onnxruntime/test/providers/cpu/tensor/affine_grid_test.cc
import argparse
import numpy as np

1019
onnxruntime/test/providers/cpu/tensor/grid_sample_test.cc Normal file
View file

File diff suppressed because it is too large Load diff

81
onnxruntime/test/providers/cpu/tensor/grid_sample_test_gen.py Normal file
View file

@ -0,0 +1,81 @@
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# This code is used to generate the test cases for the GridSample operator
# in onnxruntime/test/providers/cpu/tensor/grid_sample_test.cc
import torch
# Define the input dimensions
N, C, D, H, W = 2, 2, 3, 3, 2
# Define the modes, padding modes, and whether to align corners
modes = ["nearest", "bilinear", "bicubic"]
padding_modes = ["zeros", "border", "reflection"]
align_corners_options = [True, False]
# Loop over the combinations of parameters
torch.manual_seed(0)
for opset_version in [16, 20]:
for mode in modes:
for padding_mode in padding_modes:
for align_corners in align_corners_options:
for ndim in [4, 5]:
if ndim == 5 and mode == "bicubic":
continue
if opset_version < 20 and ndim == 5:
continue
# Create a random input tensor with the specified dimensions
input_shape = (N,) + (C,) + (((D, H, W)) if ndim == 5 else ((H, W)))
input_tensor = torch.randn(*input_shape)
# Create a random grid tensor with the specified dimensions
grid_shape = (N,) + (((D, H, W)) if ndim == 5 else ((H, W))) + (ndim - 2,)
# Between -1.2 to + 1.2
grid_tensor = torch.rand(*grid_shape) * 2.4 - 1.2
# Apply grid_sample
output_tensor = torch.nn.functional.grid_sample(
input_tensor, grid_tensor, mode=mode, padding_mode=padding_mode, align_corners=align_corners
)
X_data_str = "{" + ", ".join([f"{x:.6f}f" for x in input_tensor.numpy().flatten()]) + "}"
Grid_data_str = "{" + ", ".join([f"{x:.6f}f" for x in grid_tensor.numpy().flatten()]) + "}"
Y_shape = output_tensor.shape
Y_data_str = "{" + ", ".join([f"{x:.6f}f" for x in output_tensor.numpy().flatten()]) + "}"
onnx_mode = mode
if opset_version >= 20:
if mode == "bilinear":
onnx_mode = "linear"
elif mode == "bicubic":
onnx_mode = "cubic"
onnx_align_corners = 1 if align_corners else 0
test_name = f"test_grid_sample_{opset_version}_{ndim}D_{mode}_{padding_mode}_{'align_corners' if align_corners else 'no_align_corners'}"
print(f"TEST(GridsampleTest, {test_name}) {{")
print(f'OpTester test("GridSample", {opset_version});')
print(f'std::string mode = "{onnx_mode}";')
print(f'std::string padding_mode = "{padding_mode}";')
print(f"int64_t align_corners = {onnx_align_corners};")
print(f"std::initializer_list<int64_t> X_shape {{ {', '.join(map(str, input_shape))} }};")
print(f"std::initializer_list<float> X_data { X_data_str };")
print(f"std::initializer_list<int64_t> Grid_shape {{ {', '.join(map(str, grid_shape))} }};")
print(f"std::initializer_list<float> Grid_data { Grid_data_str };")
print(f"std::initializer_list<int64_t> Y_shape {{ {', '.join(map(str, Y_shape))} }};")
print(f"std::initializer_list<float> Y_data { Y_data_str };")
print('test.AddInput<float>("X", X_shape, X_data);')
print('test.AddInput<float>("Grid", Grid_shape, Grid_data);')
print('test.AddAttribute("mode", mode);')
print('test.AddAttribute("padding_mode", padding_mode);')
print('test.AddAttribute("align_corners", align_corners);')
print('test.AddOutput<float>("Y", Y_shape, Y_data);')
print("test.Run();")
print("}")
print("\n")

18
onnxruntime/test/testdata/onnx_backend_test_series_filters.jsonc vendored
View file

@ -239,24 +239,6 @@
"^test_ai_onnx_ml_label_encoder_string_int_no_default",
"^test_ai_onnx_ml_label_encoder_tensor_mapping",
"^test_ai_onnx_ml_label_encoder_tensor_value_only_mapping",
"^test_gridsample_aligncorners_true",
"^test_gridsample_bicubic_align_corners_0_additional_1",
"^test_gridsample_bicubic_align_corners_1_additional_1",
"^test_gridsample_bicubic",
"^test_gridsample_bilinear_align_corners_0_additional_1",
"^test_gridsample_bilinear_align_corners_1_additional_1",
"^test_gridsample_bilinear",
"^test_gridsample_border_padding",
"^test_gridsample",
"^test_gridsample_nearest_align_corners_0_additional_1",
"^test_gridsample_nearest_align_corners_1_additional_1",
"^test_gridsample_nearest",
"^test_gridsample_reflection_padding",
"^test_gridsample_volumetric_bilinear_align_corners_0",
"^test_gridsample_volumetric_bilinear_align_corners_1",
"^test_gridsample_volumetric_nearest_align_corners_0",
"^test_gridsample_volumetric_nearest_align_corners_1",
"^test_gridsample_zeros_padding",
"^test_image_decoder_decode_bmp_rgb",
"^test_image_decoder_decode_jpeg2k_rgb",
"^test_image_decoder_decode_jpeg_bgr",